"""
The Assumptions of this cell are modified to the test_sgc_input_PSTH.py but should be identical in function

synapse type = simple
synapse type = single
species = mouse
sgc = dummy SGC lacks cell body
cell connection = pyramidal



"""

import sys
import numpy as np
import pyqtgraph as pg
from neuron import h
from cnmodel.protocols import Protocol
from cnmodel import cells
from cnmodel.util import sound
from cnmodel.util import custom_init
import cnmodel.util.pynrnutilities as PU
from cnmodel import data

synapseType = "simple"  # no other options exist
species = "mouse"  # tables for other species do not yet exist


def main():
    prot = SGCInputTestPSTH(seed=1233451898)
    prot.run()
    prot.show()
    if sys.flags.interactive == 0:
        pg.QtGui.QApplication.exec_()


class SGCInputTestPSTH(Protocol):
    def __init__(self, temp=34.0, dt=0.025, seed=575982035):
        """
        :param temp: Celcius
        :param dt: resolution in ms
        :param seed: seed should be generated randomly

        Pre-determined variables that can be changed by editing the value next to them

        AMPA_gmax and n_sgc = loaded from data that is kept in cnmodel/data dir

        Defines a single tone parameters with which to test the neuron complex
        """
        # super only defines reset function
        super(SGCInputTestPSTH, self).__init__()
        #
        self.temp = temp
        self.dt = dt
        self.seed = seed
        self.simulator = "cochlea"
        # predetermined variables
        self.cell = "pyramidal"
        self.nrep = 50  # number of repetitions
        self.stimulus = "tone"
        self.Fs = 100e3  # in Hz
        self.f0 = 4000.0  # stimulus in Hz
        self.cf = 4000.0  # SGCs in Hz
        self.fMod = 100.0  # mod freq, Hz
        self.dMod = 0.0  # % mod depth, Hz
        self.dbspl = 50.0
        self.sr = 2  # set SR group

        # variables loaded from data
        AMPA_gmax, n_sgc = self.load_variables_from_data()
        self.n_sgc = int(np.round(n_sgc))

        # value needed for simple synapses
        self.AMPA_gmax = AMPA_gmax

        # the stimulation delivered is a Tone
        self.run_duration = 0.20  # in seconds
        self.pip_duration = 0.05  # in seconds
        self.pip_start = [0.1]  # in seconds
        self.f0 = 4000.0
        self.cf = 4000.0
        self.stim = sound.TonePip(
            rate=self.Fs,
            duration=self.run_duration,
            f0=self.f0,
            dbspl=self.dbspl,
            ramp_duration=2.5e-3,
            pip_duration=self.pip_duration,
            pip_start=self.pip_start,
        )

        # creates empty lists that will contain the results of each rep
        self.vms = [None for n in range(self.nrep)]
        self.synapses = [None for n in range(self.nrep)]
        self.xmtrs = [None for n in range(self.nrep)]
        self.pre_cells = [None for n in range(self.nrep)]
        self.time = [None for n in range(self.nrep)]

    def load_variables_from_data(self):
        AMPA_gmax = (
            data.get(
                "sgc_synapse", species=species, post_type=self.cell, field="AMPA_gmax"
            )[0]
            / 1e3
        )
        n_sgc = data.get(
            "convergence", species=species, post_type=self.cell, pre_type="sgc"
        )[0]
        print(n_sgc)
        return AMPA_gmax, n_sgc

    def check_assertations(self):
        assert self.cell in [
            "bushy",
            "tstellate",
            "octopus",
            "dstellate",
            "tuberculoventral",
            "pyramidal",
        ]
        assert self.stimulus == "tone"  # cases available
        assert synapseType == "simple"

    def run(self):
        self.check_assertations()
        # info based on the parameters set in the __init__ statement
        info = {
            "n_sgc": self.n_sgc,
            "gmax": self.AMPA_gmax,
            "stim": self.stim,
            "simulator": self.simulator,
            "cf": self.cf,
            "sr": self.sr,
            "seed": self.seed,
            "run_duration": self.run_duration,
            "temp": self.temp,
            "dt": self.dt,
            "init": custom_init,
        }

        # run number of trials based on nrep defined in __init__ statement
        for nr in range(self.nrep):
            info["seed"] = self.seed + 3 * self.n_sgc * nr
            res = run_trial(info)
            # res contains: {'time': time, 'vm': Vm, 'xmtr': xmtr, 'pre_cells': pre_cells, 'post_cell': post_cell}

            # unpacks the res dict returned from the run_trial() into refrenceable variables
            self.pre_cells[nr] = res["pre_cells"]
            self.time[nr] = res["time"]
            self.xmtr = {k: v.to_python() for k, v in res["xmtr"].items()}
            self.vms[nr] = res["vm"]
            self.synapses[nr] = res["synapses"]
            self.xmtrs[nr] = self.xmtr

        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")
            p1 = self.stimulus_graph()
            p2 = self.an_spikes_graph()
            p3 = self.cell_spikes_graph()
            p4 = self.voltage_graph()
            p5 = self.xmtr_graph()
            p6 = (
                self.an_psth_graph()
            )  # requires that an_spikes_graph() has been called before
            p7 = (
                self.cell_psth_graph()
            )  # requires that cell_spikes_graph() has been called before

            # links x axis
            p1.setXLink(p1)
            p2.setXLink(p1)
            p3.setXLink(p1)
            p4.setXLink(p1)
            p5.setXLink(p1)
            self.win.show()

        ############# Graph options to be included in the show() method ###################
        def stimulus_graph(self):
            p1 = self.win.addPlot(
                title="Stimulus", row=0, col=0, labels={"bottom": "T (ms)", "left": "V"}
            )
            p1.plot(
                self.stim.time * 1000, self.stim.sound, pen=pg.mkPen("k", width=0.75)
            )
            return p1

        def an_spikes_graph(self):
            p2 = self.win.addPlot(
                title="AN spikes",
                row=1,
                col=0,
                labels={"bottom": "T (ms)", "left": "AN spikes (first trial)"},
            )
            for nr in range(self.nrep):
                xan = []
                yan = []
                for k in range(len(self.pre_cells[nr])):
                    r = self.pre_cells[nr][k]._spiketrain
                    xan.extend(r)
                    yr = k + np.zeros_like(r) + 0.2
                    yan.extend(yr)
                c = pg.PlotCurveItem()
                xp = np.repeat(np.array(xan), 2)
                yp = np.repeat(np.array(yan), 2)
                yp[1::2] = yp[::2] + 0.6
                c.setData(
                    xp.flatten(),
                    yp.flatten(),
                    connect="pairs",
                    pen=pg.mkPen(pg.intColor(nr, self.nrep), hues=self.nrep, width=1.0),
                )
                self.xan = xan
                self.yan = yan
                p2.addItem(c)
            return p2

        def cell_spikes_graph(self):
            p3 = self.win.addPlot(
                title="%s Spikes" % self.cell,
                row=2,
                col=0,
                labels={"bottom": "T (ms)", "left": "Trial #"},
            )
            xcn = []
            ycn = []
            xspks = []
            for k in range(self.nrep):
                bspk = PU.findspikes(self.time[k], self.vms[k], -35.0)
                xcn.extend(bspk)
                yr = k + np.zeros_like(bspk) + 0.2
                ycn.extend(yr)
            d = pg.PlotCurveItem()
            xp = np.repeat(np.array(xcn), 2)
            yp = np.repeat(np.array(ycn), 2)
            yp[1::2] = yp[::2] + 0.6
            d.setData(
                xp.flatten(),
                yp.flatten(),
                connect="pairs",
                pen=pg.mkPen("k", width=1.5),
            )
            self.xcn = xcn
            self.ycn = ycn
            p3.addItem(d)

            return p3

        def voltage_graph(self):
            p4 = self.win.addPlot(
                title="%s Vm" % self.cell,
                row=3,
                col=0,
                labels={"bottom": "T (ms)", "left": "Vm (mV)"},
            )
            for nr in range(self.nrep):
                p4.plot(
                    self.time[nr],
                    self.vms[nr],
                    pen=pg.mkPen(pg.intColor(nr, self.nrep), hues=self.nrep, width=1.0),
                )
            return p4

        def xmtr_graph(self):
            p5 = self.win.addPlot(
                title="xmtr", row=0, col=1, labels={"bottom": "T (ms)", "left": "gSyn"}
            )
            if synapseType == "multisite":
                for nr in [0]:
                    syn = self.synapses[nr]
                    j = 0
                    for k in range(self.n_sgc):
                        synapse = syn[k]
                        for i in range(synapse.terminal.n_rzones):
                            p5.plot(
                                self.time[nr],
                                self.xmtrs[nr]["xmtr%04d" % j],
                                pen=pg.mkPen(
                                    pg.intColor(nr, self.nrep),
                                    hues=self.nrep,
                                    width=1.0,
                                ),
                            )
                            j = j + 1
            return p5

        def an_psth_graph(self):
            p6 = self.win.addPlot(
                title="AN PSTH",
                row=1,
                col=1,
                labels={"bottom": "T (ms)", "left": "Sp/ms/trial"},
            )
            bins = np.arange(0, 200, 1)
            (hist, binedges) = np.histogram(self.xan, bins)
            curve6 = p6.plot(
                binedges,
                hist,
                stepMode=True,
                fillBrush=(0, 0, 0, 255),
                brush=pg.mkBrush("k"),
                fillLevel=0,
            )
            return p6

        def cell_psth_graph(self):
            p7 = self.win.addPlot(
                title="%s PSTH" % self.cell,
                row=2,
                col=1,
                labels={"bottom": "T (ms)", "left": "Sp/ms/trial"},
            )
            bins = np.arange(0, 200, 1)
            (hist, binedges) = np.histogram(self.xcn, bins)
            curve7 = p7.plot(
                binedges,
                hist,
                stepMode=True,
                fillBrush=(0, 0, 0, 255),
                brush=pg.mkBrush("k"),
                fillLevel=0,
            )
            return p7


def run_trial(info):
    """
    This function is really the bread
    and butter of the run and determines the layout
    of the model and the stimulus is
    created and delivered by the SGC.

    Runs a single trial and returns it
    into another dictionary that is stored
    before the next trial is run.

    :param info: dict containing :'n_sgc': self.n_sgc,
                                  'gmax': self.AMPA_gmax,
                                  'stim': self.stim,
                                  'simulator': self.simulator,
                                  'cf': self.cf, 'sr': self.sr,
                                  'seed': self.seed,
                                  'run_duration': self.run_duration,
                                  'temp': self.temp,
                                  'dt': self.dt,
                                  'init': custom_init

    post cell allowed to use defaults in the case of things like pyramidal cells are guinea pigs instead of a mouse
    """
    # for the model to change the post_cell needs to be changed here otherwise the model will not be changed.
    # special consideration needs to be taken because not all cell parameters are universal
    # ex: Pyramidal does not have a species='mouse', nothing just assumes a default and not all defaults are the same
    post_cell = cells.Pyramidal.create()
    pre_cells = []
    synapses = []
    xmtr = {}

    # connects all of the SGC fibers to the post_cell using simple synapses and then generates a sound stim based on a
    # seed the number of repetitions desired
    for nsgc, sgc in enumerate(range(info["n_sgc"])):
        pre_cells.append(cells.DummySGC(cf=info["cf"], sr=info["sr"]))
        synapses.append(pre_cells[-1].connect(post_cell, type="simple"))
        synapses[-1].terminal.netcon.weight[0] = info["gmax"]

        # sets sounds stim for each of the SGC fibers independently
        pre_cells[-1].set_sound_stim(
            info["stim"], seed=info["seed"] + nsgc, simulator=info["simulator"]
        )

    # Recording each trial and returning it as a dictionary
    Vm = h.Vector()
    Vm.record(post_cell.soma(0.5)._ref_v)
    rtime = h.Vector()
    rtime.record(h._ref_t)
    h.tstop = 1e3 * info["run_duration"]  # duration of a run
    h.celsius = info["temp"]
    h.dt = info["dt"]
    post_cell.cell_initialize()
    info["init"]()
    h.t = 0.0
    h.run()
    return {
        "time": np.array(rtime),
        "vm": Vm.to_python(),
        "xmtr": xmtr,
        "pre_cells": pre_cells,
        "post_cell": post_cell,
        "synapses": synapses,
    }


if __name__ == "__main__":
    main()
