import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from scripts.interneuron_placement import create_grid_of_excitatory_neurons, plot_neural_sheet
from scripts.spatial_maps.orientation_map import OrientationMap
from pypet import Environment, cartesian_product, Trajectory, Result

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

TRAJ_NAME_ORIENTATION_MAPS = "precalculated_orientation_maps"

def generate_and_save_seeded_map(traj):
    corr_len = traj.corr_len
    seed = traj.seed
    print(corr_len,seed)
    N_E = 900

    sheet_x = 450
    sheet_y = 450

    number_of_excitatory_neurons_per_row = int(np.sqrt(N_E))

    map = OrientationMap(number_of_excitatory_neurons_per_row + 1, number_of_excitatory_neurons_per_row + 1,
                         corr_len, sheet_x, sheet_y, seed)

    loaded = False
    try:
        map.load_orientation_map()
        loaded = True
    except:
        print(
            'No map yet with {}x{} pixels and {} pixel correllation length and {} seed'.format(map.x_dim, map.y_dim, map.corr_len, map.rnd_seed))
    if not loaded:
        map.improve(10)
        traj.f_add_result('map', map.angle_grid, comment='map with {}x{} pixels and {} pixel correllation length and seed {}'.format(map.x_dim, map.y_dim, map.corr_len, map.rnd_seed))

        print('map_{}x{}_pixels_{}_corr_pix_{}_seed is saved'.format(map.x_dim, map.y_dim,map.corr_len,map.rnd_seed))

    return map.angle_grid


def create_map_frame(traj, result_list):
    print(len(result_list))
    corr_len_range = traj.par.f_get('corr_len').f_get_range()
    seed_range = traj.par.f_get('seed').f_get_range()

    corr_len_index = sorted(set(corr_len_range))
    seed_index = sorted(set(seed_range))
    map_frame = pd.DataFrame(columns=corr_len_index, index=seed_index)
    print(map_frame)
    # This frame is basically a two dimensional table that we can index with our
    # parameters

    # Now iterate over the results. The result list is a list of tuples, with the
    # run index at first position and our result at the second
    for result_tuple in result_list:
        run_idx = result_tuple[0]
        print(run_idx)
        result_map = result_tuple[1]
        corr_len_val = corr_len_range[run_idx]
        print(corr_len_range[run_idx])
        seed_val = seed_range[run_idx]
        # print(corr_len_val, seed_val, result_map)
        map_frame.loc[corr_len_val, seed_val] = corr_len_val *seed_val

    print(map_frame)

    # Finally we going to store our new firing rate table into the trajectory
    traj.f_add_result('summary.map_frames', map_frame=map_frame,
                      comment='Contains precalculated orientation maps')

def create_seeded_maps_in_scale_range():
    import multiprocessing
    import itertools

    env = Environment(trajectory=TRAJ_NAME_ORIENTATION_MAPS, filename=DATA_FOLDER,
                      file_title='orientation_map_set',
                      comment='A number of precalculated orientation maps!',
                      large_overview_tables=True,
                      overwrite_file=True,
                      multiproc=True,
                      log_folder=LOG_FOLDER,
                      ncores=3)
    traj = env.trajectory

    traj.f_add_parameter('corr_len', 200.0, comment='Correlation length')
    traj.f_add_parameter('seed', 1, comment='Seed for random')

    traj.f_explore(cartesian_product({'corr_len': [200.0], 'seed': [1]}))

    # env.add_postprocessing(create_map_frame)

    env.run(generate_and_save_seeded_map)

    env.disable_logging()

    # corr_len_range = range(0, 451, 15)
    # seed_range = range(10)
    # pool_arguments = itertools.product(corr_len_range,seed_range)
    # # print([*pool_arguments])
    #
    # pool = multiprocessing.Pool()
    # pool.starmap(generate_and_save_seeded_map,[*pool_arguments])


if __name__ == "__main__":
    create_seeded_maps_in_scale_range()
