import os
import pickle
from typing import Union

import numpy as np
from lmfit import Model
from lmfit.models import ExponentialModel
from pylibrary.tools import utility as pu

verbose = False

"""
Perform analysis of IVs pip from model_run
Provides basic analysis for:
    1. Current clamp IV
    2. Voltage clamp IV
    3. Count spikes and get FI

"""


class AnalyzeRun:
    def __init__(self, results):
        self.results = results
        self.injs = list(results.keys())
        self.nRun = len(self.injs)
        self.parseResults(results)
        self.parseStim(results)
        self.IVResult = {}

    def IV(self):
        """
        Compute the IV of the loaded dataset
        """
        self.analyzeIV(
            t=self.t,
            voltage=self.voltage,
            current=self.current,
            windows=self.tw,
            threshold=self.thr,
        )

    def VC(self):
        """
        Compute the IV of the loaded dataset
        """
        self.analyzeVC(
            t=self.t, voltage=self.voltage, current=self.current, windows=self.tw
        )

    def parseStim(self, result: dict):
        """
        parse the stimulus information in the results dictionary.
        We only need to look at the first element to get the delay and duration
        """
        try:
            site = result[self.injs[0]].stim
        except AttributeError:
            site = result[self.injs[0]]["stim"]
        self.delay = site["delay"]
        self.duration = site["dur"]
        self.tw = [self.delay, self.duration, 10.0]
        if verbose:
            print("analyze_run.py: self.tw = ", self.tw)

    def parseResults(self, result: dict):
        """
        Parse the results created by generate_run, building 2d arrays for V and I,
        and a vector for t

        Parameters
        ----------
        result : dict
            The dictionary of the results, as generated by model_run

        Returns
        -------
        Nothing
        """

        self.somasite = ["postsynapticV", "postsynapticI"]
        inj0 = self.injs[0]
        try:
            msite = self.results[inj0].monitor
        except AttributeError:
            msite = self.results[inj0]["monitor"]
        vlen = len(msite[self.somasite[1]])
        self.voltage = np.zeros((self.nRun, vlen))
        self.current = np.zeros((self.nRun, vlen))
        for j, i in enumerate(result.keys()):  # each result key is a current level...
            try:
                msite = result[self.injs[j]].monitor
            except AttributeError:
                msite = result[self.injs[j]]["monitor"]
            if j == 0:
                self.t = msite["time"][0:vlen]
            self.voltage[j, :] = msite[self.somasite[0]]
            self.current[j, :] = msite[self.somasite[1]]
            self.current[j, 0] = self.current[j, 1]  # fix first point problem..
        self.thr = -30.0  # mV

    def clean_spiketimes(
        self, spikeTimes: Union[list, np.ndarray], mindT: float = 0.7
    ) -> Union[list, np.ndarray]:
        """
        Clean up spike time array, removing all less than mindT
        Parameters
        ----------
        spikeTimes : list or numpy array (1-D)
            array of the spike times
        
        mindT : float (default : 0.7)
            minimum time between spikes, in the same units as spikeTimes
            (normally this will be in milliseconds)

        Return
        ------
        spikeTimes : list or numpy array (1-D_
            A cleaned list of the spike times where the events are at least
            mindT appart.
            Note: If there is only 1 or 0 spikes in array, just return the array
        """

        if len(spikeTimes) > 1:
            dst = np.diff(spikeTimes)
            st = np.array(spikeTimes[0])  # get first spike
            sok = np.where(dst > mindT)
            st = np.append(st, [spikeTimes[s + 1] for s in sok])
            spikeTimes = st
        return spikeTimes

    def analyzeIV(
        self,
        t: np.ndarray,
        voltage: np.ndarray,
        current: np.ndarray,
        windows: Union[list, np.ndarray],
        threshold: float,
    ):
        """
        Analyze a set of voltage records (IV), with spike threshold

        Parameters
        ----------
        t : numpy array of floats (1-D)
            time array for the voltage and current
        voltage : numpy array of floats (2-D)
            voltage traces to be analyzed. Dimension 0 is trace number,
            dimension 1 corresponds to time evolution of voltage
        current : numpy array of floats (2-D)
            current traces, corresponding to the voltage traces in V. Should
            be organized in the same way.
        windows : list
            list of [tdelay, tdur, tssw], where:
                tdelay is the delay to the start of the step.
                tdur is the duration of the step
                tssw is the duration of the steady-state window prior
                    to the end of the step
        threshold : float
            voltage threshold that will be used for spike detection.

        Returns
        -------
            a dictionary with:
            vmin
            vss
            i for vmin and vss
            spike count
            ispk
            (eventually should also include time constant measures,and adaptation ratio)
        """

        if verbose:
            print("starting analyzeIV")
        self.thr = threshold
        ntraces = np.shape(voltage)[0]
        # initialize all result arrays, lists and dicts
        vss = np.empty(ntraces)
        vmin = np.zeros(ntraces)
        vrmss = np.zeros(ntraces)
        vm = np.zeros(ntraces)
        ic = np.zeros(ntraces)
        nspikes = np.zeros(ntraces)
        ispikes = np.zeros(ntraces)
        t_minv = np.zeros(ntraces)
        fsl = []
        fisi = []
        spk = {}
        taus = {}
        tauih = {}
        xtfit = {}
        ytfit = {}
        xihfit = {}
        yihfit = {}
        dt = t[1] - t[0]
        # break down the time windows
        tss = [int((windows[1] - windows[2]) / dt), int(windows[1] / dt)]
        ts = windows[0]
        te = windows[1]

        for j in range(0, ntraces):
            if verbose:
                print("    analyzing trace: %d" % (j))
            vss[j] = np.mean(voltage[j, tss[0] : tss[1]])  # steady-state voltage
            ic[j] = np.mean(current[j, tss[0] : tss[1]])  # corresponding currents
            vm[j] = np.mean(
                voltage[j, 0 : int((ts - 1.0) / dt)]
            )  # resting potential - for 1 msec prior to step
            minV = np.argmin(voltage[j, int(ts / dt) : int((ts + te) / dt)])
            t_minv[j] = t[minV + int(ts / dt)]  # time of minimum
            vmin[j] = voltage[j, minV + int(ts / dt)]  # value of minimum
            spk[j] = pu.findspikes(
                t, voltage[j, :], self.thr, t0=ts, t1=te, dt=1.0, mode="peak"
            )
            spk[j] = self.clean_spiketimes(spk[j])
            nspikes[j] = spk[j].shape[0]  # build spike count list
            ispikes[j] = ic[j]  # currents at which spikes were detected
            if nspikes[j] >= 1:  # get FSL
                fsl.append(spk[j][0])
            else:
                fsl.append(None)
            if nspikes[j] >= 2:  # get first interspike interval
                fisi.append(spk[1] - spk[j][0])
            else:
                fisi.append(None)

            # fit the hyperpolarizing responses for Tau_m and "sag" due to Ih activation
            if (
                ic[j] < 0.0 and (t_minv[j] - ts) > 5.0 * dt
            ):  # just for hyperpolarizing pulses...
                if verbose:
                    print("    fitting trace %d" % j)
                    print("t.shape: ", t.shape)
                    print("V.shape: ", voltage[j, :].shape)
                    print("ts, vmin: ", ts, vmin)
                    print("ic[j]: ", ic[j])

                taus[j], xtfit[j], ytfit[j] = self.single_taufit(
                    t, voltage[j, :], ts, t_minv[j]
                )
                if verbose:
                    print("     calling fit")
                if (te - t_minv[j]) > 10.0 * dt:
                    tauih[j], xihfit[j], yihfit[j] = self.single_taufit(
                        t, voltage[j, :], t_minv[j], te + ts
                    )  # fit the end of the trace
                if verbose:
                    print("     completed fit")
            if verbose:
                print("   >>> completed analyzing trace %d" % j)
        if verbose:
            print("done with traces")
        if len(ic) > 1:
            icg = np.gradient(ic)
            icn = np.argwhere(ic <= 0.0)
            RinIVss = (np.gradient(vss)) / icg  # measure steady-state input resistance
            RinIVpk = (np.gradient(vmin)) / icg  # measure "peak" input resistance
            if verbose:
                print("ic: ", ic)
                print("icn: ", icn)
                print("icg: ", icg)
                print("vss: ", vss)
                print("vmin: ", vmin)
                print("RinIVss: ", RinIVss)
                print("RinIVpk: ", RinIVpk)

            if len(RinIVss[icn]) > 0:
                Rinss = np.max(RinIVss[icn])  # extract max
            else:
                Rinss = np.nan
            if len(RinIVpk[icn]) > 0:
                Rinpk = np.max(RinIVpk[icn])  # extract max
            else:
                Rinpk = np.nan
            print("RinSS, RinPeak: ", Rinss, Rinpk)
        else:
            Rinss = np.nan
            Rinpk = np.nan

        if verbose:
            print("building IVResult")
        self.IVResult = {
            "I": ic,
            "Vmin": vmin,
            "Vss": vss,
            "Vrmss": vrmss,
            "Vm": vm,
            "Tmin": np.array(t_minv),
            "Ispike": np.array(ispikes),
            "Nspike": np.array(nspikes),
            "Tspike": np.array(spk),
            "FSL": np.array(fsl),
            "FISI": np.array(fisi),
            "taus": taus,
            "tauih": tauih,
            "Rinss": Rinss,
            "Rinpk": Rinpk,
            "taufit": [xtfit, ytfit],
            "ihfit": [xihfit, yihfit],
        }

    def analyzeVC(
        self,
        t: np.ndarray,
        voltage: np.ndarray,
        current: np.ndarray,
        windows: Union[list, np.ndarray],
    ):
        """
        Analyze a set of current records (vc); return SS boltzmann fit

        Parameters
        ----------
        t : numpy array of floats (1-D)
            time array for the voltage and current
        V : numpy array of floats (2-D)
            voltage traces to be analyzed. Dimension 0 is trace number,
            dimension 1 corresponds to time evolution of voltage
        I : numpy array of floats (2-D)
            current traces, corresponding to the voltage traces in V. Should
            be organized in the same way.
        windows : list
            list of [tdelay, tdur, tssw], where:
                tdelay is the delay to the start of the step.
                tdur is the duration of the step
                tssw is the duration of the steady-state window prior
                    to the end of the step

        Returns
        -------
            a dictionary with:
            vmin
            vss
            i for vmin and vss
            spike count
            ispk
            (eventually should also include time constant measures,and adaptation ratio)
        """

        def boltzI(x, gmax, vhalf, k, E):
            return gmax * (1.0 / (np.exp((x - E) / k)))

        if verbose:
            print("starting analyzeVC")
        ntraces = np.shape(voltage)[0]
        # initialize all result arrays, lists and dicts
        vss = np.empty(ntraces)
        vmin = np.zeros(ntraces)
        vrmss = np.zeros(ntraces)
        vm = np.zeros(ntraces)
        ic = np.zeros(ntraces)

        dt = t[1] - t[0]
        # break down the time windows
        tss = [int((windows[1] - windows[2]) / dt), int(windows[1] / dt)]
        ts = windows[0]

        for j in range(0, ntraces):
            if verbose:
                print("    analyzing trace: %d" % (j))
            vss[j] = np.mean(voltage[j, tss[0] : tss[1]])  # steady-state voltage
            ic[j] = np.mean(current[j, tss[0] : tss[1]])  # corresponding currents
            vm[j] = np.mean(
                voltage[j, 0 : int((ts - 1.0) / dt)]
            )  # resting potential - for 1 msec prior to step

            if verbose:
                print("   >>> completed analyzing trace %d" % j)
        if verbose:
            print("done with traces")
        # now fit traces to variation of g = gmax * I/(V-Vr)
        gmodel = Model(boltzI)

        gmodel.set_param_hint("amax", value=3.0, min=0.0)
        gmodel.set_param_hint("vhalf", value=38.0, min=00.0, max=90.0)
        gmodel.set_param_hint("k", value=7.0, min=1.0, max=20.0, vary=False)
        gmodel.set_param_hint("E", value=-84.0, vary=False)
        gparams = gmodel.make_params()
        result = gmodel.fit(ic, params=gparams, x=vss)

        if verbose:
            print("building VCResult")
        self.VCResult = {
            "I": ic,
            "Vmin": vmin,
            "Vss": vss,
            "Vrmss": vrmss,
            "Vm": vm,
            "fit": result,
            "Boltz_gmax": result.params["gmax"],
            "Boltz_vhalf": result.params["vhalf"],
            "Boltz_slopefactor": result.params["k"],
        }

    def saveIVResult(self, name: Union[None, tuple]):
        """
        Save the result of multiple runs to disk. Results is in a dictionary,
        each element of which is a Param structure, which we then turn into
        a dictionary.
        Parameters
        ----------
        name : tuple of str (default: None)
            The file name is a tuple, where the first element is a path and the
            second is a string name that will be decorated with the data type (by appending)
        """
        fn = os.path.join(name[0], name[1] + "_IVresult.p")
        pfout = open(fn, "wb")
        pickle.dump({"IVResult": self.IVResult}, pfout)
        pfout.close()

    def single_taufit(self, x: np.ndarray, y: np.ndarray, t0: float, t1: float):
        """
        Perform single exponential fit to voltage traces

        Parameters
        ----------
        x : numpy array
            time corresponding to data in y
        y : numpy array
            voltage trace (single trace)
        t0 : float
            start time for fit (ms)
        t1 : float
            end time for fit (ms)

        Returns
        -------
        tuple of (best fit parameters as dict, x times, best fit y values)
        If fit fails, return is (None, None, None)
        """

        (cx, cy) = pu.clipdata(y, x, t0, t1, minFlag=False)
        expmodel = ExponentialModel()  # , prefix='', missing=None, name=None, **kwargs)
        expmodel.set_param_hint("decay", min=0.1, max=50.0)
        cye = np.mean(cy[-5:])
        result = expmodel.fit(cy - cye, x=cx - t0, amplitude=0.0, decay=10.0)
        if verbose:
            print(result.fit_report())  # print the result
        rbv = result.best_values
        fitr = {"a": 0, "tau": 5.0, "dc": 0.0}
        fitr["a"] = rbv["amplitude"]
        fitr["tau"] = rbv["decay"]
        return fitr, cx, result.best_fit + cye
