"""
Layout:

        if __name__==__main__ handles cmd args, instantiates test, runs it, displays results

        run_trial(): defines a model and then runs the test using preset params in hoc and return the info to the class
        SGCInputTest:
            __init__ : defines many static variables
            run(): calls delivers information to the run_trial() function and the recieved information of a single run
                back and stores it as an extensible list in the class
            show(): displays graphs depending on the graph options selected below it and displayed in a printout

"""
import os
import argparse
import time
import json
import sys
from random import randint
from cnmodel.protocols import Protocol
from cnmodel.util import sound
from cnmodel.util import custom_init
from graphs import *
from hoc_trial_run_dendrites import run


class NetworkSimulation(Protocol):
    """
    Tests a Single cell with input recieved from the SGC

    __init__: almost all parameters can be modified
    run(): simply loops over the run_trial() function and stores the results just
    show(): constructs the graphs using other functions

    """

    def __init__(
        self,
        temp=34.0,
        seed=randint(2500, 400000),
        nrep=10,
        stimulus="tone",
        simulator="cochlea",
        debug=True,
        dt=0.025
    ):
        """
        :param temp: (float) must be at 34 for default pyramidal cells
        :param dt: (float) determine hoc resolution
        :param seed: (int) contributes to randomization, needs to be changed to see different results (reduce this
                            number if you keep getting a timeout error
        :param nrep: (int) number of presentations !!must be changed in the __name__ function if not calling from cmd line!!
        :param stimulus: (str) must be 'tone'
        :param simulator: (str) currently using cochlea instead of matlab
        :param n_sgc:(int) This is the number of SGC fibers that connect to the post synaptic cell
        :param debug: (bool) controls most of the terminal printouts is on by default
        :param cell: (str) cell type !!must be changed in the __name__ function if not calling from cmd line!!
        """
        super().__init__()
        #
        self.debug = debug
        self.nrep = nrep
        self.stimulus = stimulus
        self.run_duration = 0.3  # in seconds
        # stim parameters
        self.pip_duration = 0.05  # in seconds
        self.pip_start = [0.14]  # in seconds
        self.Fs = 100e3  # in Hz
        self.f0 = 14013.0  # stimulus in Hz
        self.cf = 14013.0  # SGCs in Hz
        self.fMod = 100.0  # mod freq, Hz
        self.dMod = 0.0  # % mod depth, Hz
        self.dbspl = 55.0
        # usually cochlea
        self.simulator = simulator
        self.sr = 2  # set SR group
        self.seed = seed
        # physiological parameters
        self.temp = temp
        self.dt = dt
        self.synapse_type = "multisite"
        self.species = "rat"

        if self.stimulus == "tone":
            self.stim = sound.TonePip(
                rate=self.Fs,
                duration=self.run_duration,
                f0=self.f0,
                dbspl=self.dbspl,
                ramp_duration=5e-3,
                pip_duration=self.pip_duration,
                pip_start=self.pip_start,
            )
        # result containers
        self.vms = []
        self.vmtbs = []
        self.vmdss = []
        self.vmcar = []
        self.synapses = []
        self.auditory_nerve_cells = []
        self.time = []
        # debug function reports a print out of various information about the run
        if self.debug:
            print("Small Network Test Created")
            print("#" * 70)
            print(f"Running Test of Network with Simulated SGC fibers")
            print()
            print(f"Run Conditions:         Run Time: {self.run_duration}s,")
            print(f"                        Run Temp: {self.temp} ")
            print(f"                        Number of Presentations: {nrep}")
            print()
            print(f"Stimulus Conditions:    Type: {stimulus}")
            print(f"                        Stim Duration: {self.pip_duration}s")
            print(f"                        Characteristic F: {self.cf}hz")
            print(f"                        Stim Start:{str(self.pip_start)}s")

    def run(self, processes=1, **kwargs):
        """
        Runs the trials with the number of nreps set for the trial, calls a multiprocess run of however many trials
        are needed.
        If processes set to one then the it will run without the multiprocessing library(more reliable)
        """
        super().run()
        info = {
            "stim": self.stim,
            "simulator": self.simulator,
            "cf": self.cf,
            "sr": self.sr,
            "run_duration": self.run_duration,
            "synapse_type": self.synapse_type,
            "temp": self.temp,
            "dt": self.dt,
            "init": custom_init,
        }
        # Generate inputs
        b = [(self.seed + randint(0, 80000)) for _ in range(self.nrep)]
        c = [info for _ in range(self.nrep)]
        d = range(1, self.nrep+1)

        run_input = zip(b, c, d)
        self.all_results = run(run_input, processes=processes)
        self.unpack_data(self.all_results)

    def unpack_data(self, run_data):
        len_of_data = 0
        for nr, res in enumerate(run_data):
            try:
                len_of_data += 1
                # res contains: {'time': time, 'vm': list(Vm), 'auditory_nerve_cells': auditory_nerve_cells._spiketrain,'vmtb': list(Vmtb)}
                self.auditory_nerve_cells.append(res["auditory_nerve_cells"])
                self.time.append(res["time"])
                self.vms.append(res["vm"])
                self.vmtbs.append(res["vmtb"])
                self.vmdss.append(res["vmds"])
                self.vmcar.append(res["vmcar"])
            except(KeyError):
                continue
        self.nrep = len_of_data

    def show(self):
        """
        Creates a single page graph that contains all of the graphs based on the graphical functions in the class
        """
        self.win = pg.GraphicsWindow()
        self.win.setBackground("w")
        p2 = an_spike_graph(self.win, self.auditory_nerve_cells, 0, 0)
        p3 = spike_graph(self.win, self.time, self.vms, 1, 0, title="Pyramidal")
        p5 = spike_graph(
            self.win, self.time, self.vmtbs, 2, 0, title="Tuberculoventral"
        )
        p6 = spike_graph(self.win, self.time, self.vmdss, 3, 0, title="D-Stellate")
        p1 = stimulus_graph(self.win, self.stim, 0, 1)
        p4 = voltage_graph(self.win, self.time, self.vms, 1, 1)
        p7 = an_psth_graph(
            self.win, self.auditory_nerve_cells, 2, 1
        )
        p8 = cell_psth_graph(
            self.win, self.time, self.vms, 3, 1, title="Pyramidal"
        )

        # links x axis
        p1.setXLink(p1)
        p2.setXLink(p1)
        p3.setXLink(p1)
        p4.setXLink(p1)
        p5.setXLink(p1)
        p6.setXLink(p1)
        p7.setXLink(p1)
        p8.setXLink(p1)

        self.win.show()
        if self.debug:
            print("finished")

    def export(self):
        if self.debug:
            print("Exporting File Binary")
        t = time.gmtime()
        destination_name = f"{os.path.basename(__file__).strip('.py')}_{t.tm_mday}-{t.tm_mon}-{t.tm_year}_{t.tm_hour}_{t.tm_min}.json"
        os.scandir()
        dirname = os.path.join(os.path.dirname(__file__), "run_data")
        if not os.path.isdir(dirname):
            os.mkdir(dirname)
        filepath = os.path.join(dirname, destination_name)
        with open(filepath, "w") as f:
            json.dump(self.all_results, f)
        if self.debug:
            print(f"Run saved to {filepath}")

    def load(self, load_file=None):
        with open(load_file, "r") as file_in:
            data_list = json.load(file_in)
        self.unpack_data(data_list)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Compute Neuron response for small network of DCN"
    )
    # parser.add_argument(type=str, dest='cell', default='pyramidal',
    #                     choices=['bushy', 'tstellate', 'dstellate', 'octopus',
    #                              'tuberculoventral', 'pyramidal'],
    #                     help='Select target cell')
    parser.add_argument(
        "-n",
        "--nrep",
        type=int,
        dest="nrep",
        default=4,
        help="Set number of repetitions",
    )
    parser.add_argument(
        "-l",
        "--load",
        type=str,
        dest="load_file",
        default=None,
        help="Load data from a previous run pickle",
    )
    parser.add_argument(
        "-s",
        "--save",
        type=str,
        dest="save_flag",
        default=True,
        help="If you do not want to export file set flag to False",
    )
    parser.add_argument(
        "-d",
        "--debug",
        type=str,
        dest="debug_flag",
        default=True,
        help="If you do not want to see debug text or GUI set to False",
    )
    parser.add_argument(
        "-p",
        "--process",
        type=int,
        dest="processes",
        default=1,
        help="Set the number of processes that the run will be run across",
    )
    args = parser.parse_args()
    start = time.time()
    nrep = args.nrep
    processes = args.processes
    load_file = args.load_file

    # manages how options affect the class that is created to manage the trial run
    if args.save_flag == "False":
        save = False
    else:
        save = True
    if args.debug_flag == "False":
        debug_flag = False
    else:
        debug_flag = True

    if not load_file:
        test = NetworkSimulation(nrep=nrep, debug=debug_flag)
        test.run(processes=processes)
        if debug_flag:
            test.show()
        if save:
            test.export()
    else:
        if os.path.exists(load_file):
            test = NetworkSimulation(nrep=nrep, debug=False)
            test.load(load_file)
            test.show()
        else:
            raise FileNotFoundError(f"{load_file} does not exist")
    dtime = time.time() - start
    print("#" * 70)
    print(f"Total Elapsed Time {dtime/60} mins")

    if sys.flags.interactive == 0:
        pg.QtGui.QApplication.exec_()
