import json
import os
import pickle
import pprint

import numpy as np


def load_json(path_to_json):
    with open(path_to_json) as json_file:
        dictionary = json.load(json_file)
    print("Loaded {:s} \n{:s}".format(path_to_json, pprint.pformat(dictionary)))
    return dictionary


def save_to_pickle(path_to_pickle, result_dict):
    with open(path_to_pickle, mode="wb") as result_file:
        pickle.dump(result_dict, result_file)
    print("Save {:s}".format(path_to_pickle))

def load_pickle(path_to_pickle):
    with open(path_to_pickle, 'rb') as file:
        unpickled = pickle.load(file)
    print("Load {:s}".format(path_to_pickle))
    return unpickled


class IO:
    def __init__(self, path_to_simulation_dir):
        self.path_to_simulation_dir = path_to_simulation_dir
        self.output_folder = os.path.join(path_to_simulation_dir, "out")

    def _create_output_if_it_does_not_exist(self):
        os.makedirs(self.output_folder, exist_ok=True)

    def _create_figure_folder_if_it_does_not_exist(self):
        os.makedirs(os.path.join(self.output_folder, "figures"), exist_ok=True)

    def _load_simulation_configuration(self, filename):
        path_to_json = os.path.join(self.path_to_simulation_dir, "{:s}.json".format(filename))
        dictionary = load_json(path_to_json)
        return dictionary

    def _save_results_dict(self, filename, result_dict):
        path_to_pickle = os.path.join(self.output_folder, "{:s}.pickle".format(filename))
        save_to_pickle(path_to_pickle, result_dict)

    def _load_results_dict(self, result_name):
        path_to_pickle = os.path.join(self.output_folder, "{:s}.pickle".format(result_name))
        return load_pickle(path_to_pickle)


    def load_morphology_dict(self):
        return self._load_simulation_configuration("morphology")

    def load_dynamics_dict(self):
        return self._load_simulation_configuration("dynamics")

    def load_experiment(self):
        return self._load_simulation_configuration("experiment")

    def save_positions_and_orientations(self, excitatory_somata_positions, inhibitory_somata_positions,
                                        inhibitory_axonal_orientations):
        self._create_output_if_it_does_not_exist()
        save_array(self.output_folder, "ex_positions", excitatory_somata_positions)
        save_array(self.output_folder, "in_positions", inhibitory_somata_positions)
        save_array(self.output_folder, "in_orientations", inhibitory_axonal_orientations)

    def load_positions_and_orientations(self):
        excitatory_somata_positions = load_array(self.output_folder, "ex_positions")
        inhibitory_somata_positions = load_array(self.output_folder, "in_positions")
        inhibitory_axonal_orientations = load_array(self.output_folder, "in_orientations")
        return excitatory_somata_positions, inhibitory_somata_positions, inhibitory_axonal_orientations

    def save_in_ex_connectivity(self, in_ex_connectivity):
        save_array(self.output_folder, "in_ex_connectivity", in_ex_connectivity)

    def load_in_ex_connectivity(self):
        return load_array(self.output_folder, "in_ex_connectivity")

    def save_spike_trains(self, excitatory_spikes, inhibitory_spikes):
        self._save_results_dict("excitatory_spikes", excitatory_spikes)
        self._save_results_dict("inhibitory_spikes", inhibitory_spikes)

    def load_spike_trains(self):
        excitatory_spikes = self._load_results_dict("excitatory_spikes")
        inhibitory_spikes = self._load_results_dict("inhibitory_spikes")
        return excitatory_spikes, inhibitory_spikes

    def save_figure(self, fig, filename):
        self._create_figure_folder_if_it_does_not_exist()
        full_file_name = os.path.join(self.output_folder, "figures", "{:s}.png".format(filename))
        print("Save {:s}".format(full_file_name))
        fig.savefig(full_file_name, dpi=300)


def save_array(dir, filename, array):
    full_path = os.path.join(dir, "{:s}.npy".format(filename))
    print("Save {:s}".format(full_path))
    np.save(full_path, array)


def load_array(dir, filename):
    full_path = os.path.join(dir, "{:s}.npy".format(filename))
    print("Load {:s}".format(full_path))
    return np.load(full_path)


def get_slice_dimensions(morphology_dict):
    slice = morphology_dict["slice"]
    v_x = slice["x"]
    v_y = slice["y"]
    v_z = slice["z"]
    return v_x, v_y, v_z
