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 tqdm import tqdm

from scripts.interneuron_placement import create_grid_of_excitatory_neurons, \
    create_interneuron_sheet_by_repulsive_force

use_saved_array = False

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"

corr_len_list = range(0,451,15) #For these values maps exist
ellipse_trial_entropy_list = []
circle_trial_entropy_list = []

if not use_saved_array:
    for corr_len in tqdm(corr_len_list, desc="Calculating entropy over scale"): #TODO: Add trials, since the maps are random
        ellipse_single_trial_entropy_list = []
        circle_single_trial_entropy_list = []
        for seed in range(10):
            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))
                    continue
                tuning_map = lambda x, y: map.tuning(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)
                ellipse_single_trial_entropy_list.append(ellipse_single_trial_entropy)
                circle_single_trial_entropy_list.append(circle_single_trial_entropy)
        ellipse_trial_entropy_list.append(ellipse_single_trial_entropy_list)
        circle_trial_entropy_list.append(circle_single_trial_entropy_list)

        # interneuron_tunings = [inhibitory_axonal_clouds, inhibitory_axonal_circles]
        # plot_neural_sheet(ex_positions, ex_tunings, inhibitory_axonal_clouds)
    np.save('../../simulations/2020_02_27_entropy_over_noise_scale/circle_trial_entropy_list.npy', circle_trial_entropy_list)
    np.save('../../simulations/2020_02_27_entropy_over_noise_scale/ellipse_trial_entropy_list.npy', ellipse_trial_entropy_list)
else:
    circle_trial_entropy_list = np.load(
        '../../simulations/2020_02_27_entropy_over_noise_scale/circle_trial_entropy_list.npy')
    ellipse_trial_entropy_list = np.load(
        '../../simulations/2020_02_27_entropy_over_noise_scale/ellipse_trial_entropy_list.npy')

ellipse_entropy_mean = np.array([np.mean(i) for i in ellipse_trial_entropy_list])
circle_entropy_mean = np.array([np.mean(i) for i in circle_trial_entropy_list])

ellipse_entropy_std_dev = np.array([np.std(i) for i in ellipse_trial_entropy_list])
circle_entropy_std_dev = np.array([np.std(i) for i in circle_trial_entropy_list])
print(ellipse_trial_entropy_list)
print(ellipse_entropy_std_dev)

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