import matplotlib.pyplot as plt
import noise
import numpy as np
from brian2.units import *
from scripts.spatial_maps.orientation_map import OrientationMap
from scripts.interneuron_placement import create_interneuron_sheet_entropy_max_orientation

from scripts.interneuron_placement import create_grid_of_excitatory_neurons, \
    create_interneuron_sheet_by_repulsive_force

import multiprocessing
import itertools



trials_per_scale = 1

N_E = 900
N_I = 90

sheet_x = 450 * um
sheet_y = 450 * um

inhibitory_axon_long_axis = 100 * um
inhibitory_axon_short_axis = 25 * um

number_of_excitatory_neurons_per_row = int(np.sqrt(N_E))

'''
Tuning Maps
'''

# tuning_label = "Perlin"
tuning_label = "Orientation map"

# optimization_label = "Repulsive"
optimization_label = "Entropy Optimization"

ellipse_trial_sharpening_list = []
circle_trial_sharpening_list = []
no_conn_trial_sharpening_list = []

def get_fwhm_for_corr_len_and_seed(corr_len,seed):
    print(corr_len, seed)
    if tuning_label == "Perlin":  # TODO: How to handle scale in Perlin
        tuning_map = lambda x, y: noise.pnoise2(x / 100.0, y / 100.0, octaves=2) * np.pi
    elif tuning_label == "Orientation map":
        map = OrientationMap(number_of_excitatory_neurons_per_row + 1, number_of_excitatory_neurons_per_row + 1,
                             corr_len, sheet_x / um, sheet_y / um, seed)
        # map.improve(10)
        try:
            map.load_orientation_map()
        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))
            return -1,-1,-1
        tuning_map = lambda x, y: map.orientation_map(x, y)

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

    inhibitory_radial_axis = np.sqrt(inhibitory_axon_long_axis * inhibitory_axon_short_axis)

    if optimization_label == "Repulsive":
        inhibitory_axonal_clouds = create_interneuron_sheet_by_repulsive_force(N_I, inhibitory_axon_long_axis / um,
                                                            inhibitory_axon_short_axis / um, sheet_x / um,
                                                            sheet_y / um, random_seed=2, n_iterations=1000)

        inhibitory_axonal_circles = create_interneuron_sheet_by_repulsive_force(N_I, inhibitory_radial_axis / um,
                                                             inhibitory_radial_axis / um, sheet_x / um,
                                                             sheet_y / um, random_seed=2, n_iterations=1000)
    elif optimization_label == "Entropy Optimization":
        inhibitory_axonal_clouds, ellipse_single_trial_entropy = create_interneuron_sheet_entropy_max_orientation(
            ex_positions, ex_tunings, N_I, inhibitory_axon_long_axis / um,
            inhibitory_axon_short_axis / um, sheet_x / um,
            sheet_y / um, trial_orientations=30)
        inhibitory_axonal_circles, circle_single_trial_entropy = create_interneuron_sheet_entropy_max_orientation(
            ex_positions, ex_tunings, N_I, inhibitory_radial_axis / um,
            inhibitory_radial_axis / um, sheet_x / um,
            sheet_y / um, trial_orientations=1)

    return [ellipse_single_trial_entropy,circle_single_trial_entropy]

corr_len_range = range(0, 451, 15)
corr_len_range_len = len(corr_len_range)
seed_range = range(10)
seed_range_len = len(seed_range)
pool_arguments = itertools.product(corr_len_range,seed_range)

use_saved_array = False
if not use_saved_array:

    pool = multiprocessing.Pool()
    data = pool.starmap(get_fwhm_for_corr_len_and_seed,[*pool_arguments])
    print(type(data))
    print(data)
    data_array = np.reshape(np.array(data),(corr_len_range_len,seed_range_len,2))
    data_test = []

    np.save('../../simulations/2020_02_27_entropy_over_noise_scale/data_test_save.npy', data_array)
else:
    data_array = np.load('../../simulations/2020_02_27_entropy_over_noise_scale/data_test_save.npy')
# ellipse_trial_sharpening_list = np.zeros((len(corr_len_range),len(seed_range)))
# circle_trial_sharpening_list = np.zeros((len(corr_len_range),len(seed_range)))
# no_conn_trial_sharpening_list = np.zeros((len(corr_len_range),len(seed_range)))
print(data_array)
no_conn_trial_sharpening_list = []
ellipse_trial_sharpening_list = []
circle_trial_sharpening_list = []

for i in range(corr_len_range_len):
    ellipse_trial_sharpening_list.append(data_array[i,:,0])
    circle_trial_sharpening_list.append(data_array[i,:,1])
print(circle_trial_sharpening_list)


ellipse_sharpening_mean = np.array([np.mean(i) for i in ellipse_trial_sharpening_list])
circle_sharpening_mean = np.array([np.mean(i) for i in circle_trial_sharpening_list])

ellipse_sharpening_std_dev = np.array([np.std(i) for i in ellipse_trial_sharpening_list])
circle_sharpening_std_dev = np.array([np.std(i) for i in circle_trial_sharpening_list])
# print(ellipse_trial_sharpening_list)
# print(ellipse_entropy_std_dev)

plt.figure()
plt.plot(corr_len_range,circle_sharpening_mean, label='Circle', marker='o',color='C1')
plt.fill_between(corr_len_range,circle_sharpening_mean-circle_sharpening_std_dev,circle_sharpening_mean+circle_sharpening_std_dev,color='C1',alpha=0.4)
plt.plot(corr_len_range,ellipse_sharpening_mean, label='Ellipse', marker='o',color='C2')
plt.fill_between(corr_len_range,ellipse_sharpening_mean-ellipse_sharpening_std_dev,ellipse_sharpening_mean+ellipse_sharpening_std_dev,color='C2',alpha=0.4)
plt.xlabel('Correlation length')
plt.ylabel('Entropy')
plt.legend()
plt.show()


# plt.figure()
# plt.plot(corr_len_list,circle_entropy_mean, label='Circle', marker='o',color='C1')
# plt.fill_between(corr_len_list,circle_entropy_mean-circle_entropy_std_dev,circle_entropy_mean+circle_entropy_std_dev,color='C1',alpha=0.4)
# plt.plot(corr_len_list,ellipse_entropy_mean, label='Ellipse', marker='o',color='C2')
# plt.fill_between(corr_len_list,ellipse_entropy_mean-ellipse_entropy_std_dev,ellipse_entropy_mean+ellipse_entropy_std_dev,color='C2',alpha=0.4)
# plt.xlabel('Correlation length')
# plt.ylabel('Entropy')
# plt.legend()
# plt.show()