import numpy as np
import nixio as nix
import os
from numpy.core.fromnumeric import repeat
from traitlets.traitlets import Instance
from chirp_ams import get_signals, other_freq, other_signal, self_freq
from model import simulate, load_models
from IPython import embed
import matplotlib.pyplot as plt

def append_settings(section, sec_name, sec_type, settings):
    section = section.create_section(sec_name, sec_type)
    for k in settings.keys():
        if isinstance(settings[k], dict):
            append_settings(section, k, "settings", settings[k])
        else:
            if isinstance(settings[k], np.ndarray):
                if len(settings[k].shape) == 1:
                    section[k] = list(settings[k])
            else:
                section[k] = settings[k]

    
def save(filename, name, stimulus_settings, model_settings, self_signal, other_signal, self_freq, other_freq, complete_stimulus, responses, overwrite=False):
    if os.path.exists(filename) and not overwrite:
        nf = nix.File.open(filename, nix.FileMode.ReadWrite)
    else:
        nf = nix.File.open(filename, mode=nix.FileMode.Overwrite,
                               compression=nix.Compression.DeflateNormal)
    
    if name in nf.blocks:
        print("Data with this name is already stored! ", name)
        nf.close()
        return
    mdata = nf.create_section(name, "nix.simulation")
    append_settings(mdata, "model parameter", "nix.model.settings", model_settings)
    append_settings(mdata, "stimulus parameter", "nix.stimulus.settings", stimulus_settings)
    b = nf.create_block(name, "nix.simulation")
    b.metadata = mdata
    
    # save stimulus
    stim_da = b.create_data_array("complete_stimulus", "nix.timeseries.sampled", dtype=nix.DataType.Float,
                                  data=complete_stimulus)
    stim_da.label = "voltage"
    stim_da.label = "mv/cm"
    dim = stim_da.append_sampled_dimension(model_settings["deltat"])
    dim.label = "time"
    dim.unit = "s"
    
    # save responses
    for i in range(len(responses)):
        da = b.create_data_array("response_%i" %i, "nix.timeseries.events.spike_times", 
                                 dtype=nix.DataType.Float, data=responses[i])
        da.label = "time"
        da.unit = "s"
        dim = da.append_range_dimension()
        dim.link_data_array(da, [-1])
        
    
    nf.close()


def simulate_responses(stimulus_params, model_params, repeats=10, deltaf=20):
    cell_params = model_params.copy()
    cell = cell_params["cell"]
    del cell_params["cell"]
    del cell_params["EODf"]
    for c in stimulus_params["contrasts"]:
        print("\t\tcreating stimuli ... ")
        params = stimulus_params.copy()
        del params["contrasts"]
        del params["chirp_frequency"]
        params["contrast"] = c
        time, self_signal, self_freq, other_signal, other_freq = get_signals(**params)
        signal = (self_signal + other_signal)
        signal /= np.max(signal)
        print("\t\tcreating p-unit responses ...")
        
        spikes = []
        no_other_spikes = []
        for r in range(repeats):
            spikes.append(simulate(signal, **cell_params))
            no_other_spikes.append(simulate, self_signal))
        print("\t\tsaving ...")
        save("cell_%s.nix" % cell, "contrast_%.3f_condition_%s_deltaf_%i" %(c, stimulus_params["condition"], deltaf), params,
             cell_params, self_signal, other_signal, self_freq, other_freq, signal, spikes)


def main():
    models = load_models("models.csv")
    deltafs = [-200, -100, -20, 20, 100, 200]  # Hz, difference frequency between self and other
    
    for cell_id in range(len(models)):
        model_params = models[cell_id]
        print("Cell: %s" % model_params["cell"])
        for deltaf in deltafs:
            print("\t Deltaf: %i" % deltaf)
            stimulus_params = {"eodfs": {"self": model_params["EODf"],
                                        "other": model_params["EODf"] + deltaf},
                            "contrasts": [20, 10, 5, 2.5, 1.25, 0.625, 0.3125],
                            "chirp_size": 100,  # Hz, frequency excursion
                            "chirp_duration": 0.015,  # s, chirp duration
                            "chirp_amplitude_dip": 0.05,  # %, amplitude drop during chirp
                            "chirp_frequency": 5,  # Hz, how often does the fish chirp
                            "duration": 5.,  # s, total duration of simulation
                            "dt": model_params["deltat"],  # s, stepsize of the simulation
            }
            chirp_times = np.arange(stimulus_params["chirp_duration"],
                                    stimulus_params["duration"] - stimulus_params["chirp_duration"], 
                                    1./stimulus_params["chirp_frequency"])
            stimulus_params["chirp_times"] = chirp_times
            conditions = ["other", "self"]
        
            for c in conditions:
                print("\t\tcondition: %s" % c)
                stimulus_params["condition"] = c
                simulate_responses(stimulus_params, model_params, repeats=25,
                                   deltaf=deltaf)
    
if __name__ == "__main__":
    main()