import matplotlib.pyplot as plt
import numpy as np
from numpy.linalg import norm
from scripts.interneuron_placement import create_grid_of_excitatory_neurons, \
        plot_neural_sheet
from scripts.spatial_maps.orientation_map import OrientationMap

def plotting(X, Y, Z, o_map, ax=None, field=True):
    for y_idx in range(Z.shape[1]):
        for x_idx in range(Z.shape[0]):
            o_map_val = o_map.orientation_map(x_idx, y_idx)
            Z[x_idx, y_idx] = o_map_val
    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(111)
    plt.set_cmap('twilight')
    pcm = ax.pcolormesh(X, Y, Z.T)
    if field is True:
        for y_idx in range(Z.shape[1]):
            for x_idx in range(Z.shape[0]):
                vec = o_map.vec_grid[x_idx, y_idx]
                ax.quiver(x_idx, y_idx, vec[0], vec[1], scale=0.02, units="x", scale_units="x", width=0.09)
    plt.gcf().colorbar(pcm, ax=ax)
    # plt.show()


def create_maps_in_scale_range():



    N_E = 900

    sheet_x = 450
    sheet_y = 450

    number_of_excitatory_neurons_per_row = int(np.sqrt(N_E))

    for corr_len in range(0,301,15):
        map = OrientationMap(number_of_excitatory_neurons_per_row +1, number_of_excitatory_neurons_per_row +1,
                             corr_len, sheet_x, sheet_y)

        loaded = False
        try:
            map.load_orientation_map()
            loaded = True
        except:
            print('No map yet with {}x{} pixels and {} pixel correllation length'.format(map.x_dim,map.y_dim,map.corr_len))
        if not loaded:
            map.improve(10)
            np.save('../noise_maps/map_{}x{}_pixels_{}_corr_pix.npy'.format(map.x_dim,map.y_dim,map.corr_len),map.angle_grid)


    tuning_map = lambda x, y: map.orientation_map(x, y)

    ex_positions, ex_tunings = create_grid_of_excitatory_neurons(sheet_x, sheet_y,
                                                                 number_of_excitatory_neurons_per_row, tuning_map)

    plot_neural_sheet(ex_positions, ex_tunings)
    plt.figure()
    plt.hist(ex_tunings)

    plt.hist(Z.flat, 30)
    plt.show()

def generate_and_save_seeded_map(corr_len,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)
        np.save(
            '../noise_maps/seeded_orientation_maps/map_{}x{}_pixels_{}_corr_pix_{}_seed.npy'.format(map.x_dim,
                                                                                                    map.y_dim,
                                                                                                    map.corr_len,
                                                                                                    map.rnd_seed),
            map.angle_grid)

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


def create_seeded_maps_in_scale_range():
    import multiprocessing
    import itertools

    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])

def test():

    N_E = 900

    sheet_x = 450
    sheet_y = 450

    corr_len = 150

    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)

    map.improve(5)

    tuning_map = lambda x, y: map.orientation_map(x, y)

    ex_positions, ex_tunings = create_grid_of_excitatory_neurons(sheet_x, sheet_y,
                                                                 number_of_excitatory_neurons_per_row, tuning_map)

    plot_neural_sheet(ex_positions, ex_tunings)
    plt.figure()
    plt.hist(ex_tunings)

    plt.show()


if __name__ == "__main__":
    # create_maps_in_scale_range()
    # test()
    create_seeded_maps_in_scale_range()
