import h5py
import numpy as np

import scripts.spatial_layout
from scripts.interneuron_placement import get_optimal_orientations_for_maximum_entropy
from scripts.spatial_network.run_entropy_maximisation_orientation_map import get_orientation_map
from scripts.spatial_packing_entropy.sample_spatial_layout import correlation_length, map_seed, sheet_size, NE, \
    data_folder, NI, long_axis, short_axis, number_of_bins, ds

if __name__ == "__main__":

    orientation_map = get_orientation_map(correlation_length, map_seed, sheet_size, NE, data_folder=data_folder)

    spatial_layout = scripts.spatial_layout.SpatialLayout(orientation_map, NE, NI, long_axis, short_axis, sheet_size)

    _, optimal_orientations = get_optimal_orientations_for_maximum_entropy(spatial_layout.ex_positions,
                                                                  spatial_layout.ex_tunings,
                                                 spatial_layout.in_positions, spatial_layout.long_axis,
                                                 spatial_layout.short_axis, number_of_bins, trial_orientations=50)
    optimal_orientations = np.array(optimal_orientations)

    with h5py.File(data_folder + 'spatial_layouts.hdf5', "a") as f:
        entropy = 'max_entropy'
        if entropy not in f.keys():
            entropy_grp = f.create_group(entropy)
        else:
            entropy_grp = f[entropy]

        entropy_grp.create_dataset("entropy", shape=(), data=spatial_layout.get_mean_entropy(optimal_orientations, number_of_bins))
        entropy_grp.create_dataset("overlap", shape=(), data=spatial_layout.get_overlap(optimal_orientations, ds))
        entropy_grp.create_dataset("orientations", shape=(NI,), data = optimal_orientations)