from scripts.spatial_maps.supplement_pinwheel_map.pinwheel_map import PinwheelMap
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning) # Otherwise pypet's FutureWarning spams the console output
from pypet import Environment, cartesian_product
from scripts.spatial_network.supplement_pinwheel_map.run_simulation_pinwheel_map import SCALE_RANGE, SEED_RANGE, \
    TRAJ_NAME_PINWHEEL_MAPS

DATA_FOLDER = "../../../data/"
LOG_FOLDER = "../../../logs/"


def generate_and_save_seeded_map(traj):
    scale = traj.scale
    seed = traj.seed
    print('Generating pinwheel map with scale {} and seed {}'.format(scale, seed))

    dim = 60
    size = 900

    pinwheel_map = PinwheelMap(dim, dim, scale, size, size, seed)

    pinwheel_map.improve(10)
    traj.f_add_result('pinwheel_map', pinwheel_map.angle_grid, comment='Pinwheel map with scale {} and seed {}'.format(pinwheel_map.scale, pinwheel_map.seed))

    print('Pinwheel map with scale {} and seed {} saved'.format(pinwheel_map.scale, pinwheel_map.seed))

    return pinwheel_map.angle_grid


def create_seeded_maps_in_scale_range():
    env = Environment(trajectory=TRAJ_NAME_PINWHEEL_MAPS, filename=DATA_FOLDER,
                      file_title='pinwheel_map_set',
                      comment='A number of precalculated pinwheel maps!',
                      large_overview_tables=True,
                      overwrite_file=True,
                      multiproc=True,
                      log_folder=LOG_FOLDER,
                      ncores=0)
    traj = env.trajectory

    traj.f_add_parameter('scale', 200.0, comment='Scale factor c, see methods')
    traj.f_add_parameter('seed', 1, comment='Seed for random initialization')

    scale_range = SCALE_RANGE
    seed_range = SEED_RANGE

    traj.f_explore(cartesian_product({'scale': scale_range, 'seed': seed_range}))

    env.run(generate_and_save_seeded_map)

    env.disable_logging()


if __name__ == "__main__":
    create_seeded_maps_in_scale_range()