
import pdb
import math
import warnings
import random
import pandas as pd
import numpy as np
import torch
import matplotlib.pyplot as plt
from .sampler import Sampler

def timeRescaling(Y, pk):
    n = int(torch.sum(Y).item())
    rISIs = torch.zeros(n-1, dtype=torch.double)
    spikeindicies = Y.nonzero()
    for r in range(n-1):
        ind1 = spikeindicies[r].item()
        ind2 = spikeindicies[r+1].item()
        rISIs[r] = torch.sum(pk[ind1+1:ind2+1]).item()
    return rISIs

def timeRescalingForUnbinnedSpikes(spikesTimes, cifValues, t0, tf, dt):
    pk = cifValues*dt
    bins = torch.arange(t0-dt/2, tf+dt/2, dt)
    # start binning spikes using pandas
    cutRes, _ = pd.cut(spikesTimes, bins=bins, retbins=True)
    Y = torch.from_numpy(cutRes.value_counts().values)
    # end binning spikes using pandas
    indicesMoreThanOneSpikes = (Y>1).nonzero()
    if len(indicesMoreThanOneSpikes)>0:
        warnings.warn("Found more than one spike in {:d} bins".format(len(indicesMoreThanOneSpikes)))
        Y[indicesMoreThanOneSpikes] = 1.0
    zExp = timeRescaling(Y=Y, pk=pk)
    return zExp

def KSTestTimeRescalingNumericalCorrection(spikesTimes, cifTimes, cifValues, gamma):
    # from Haslinger et al., 2010, p. 2492,
    # Procedure for Numerical Correction

    def subtractECDFs(ecdf1X, ecdf1Y, ecdf2X, ecdf2Y):
        subX = torch.sort(torch.cat((ecdf1X, ecdf2X)))[0].numpy()
        ecdf1YI = np.interp(subX, ecdf1X.numpy(), ecdf1Y.numpy())
        ecdf2YI = np.interp(subX, ecdf2X.numpy(), ecdf2Y.numpy())
        subY = ecdf1YI-ecdf2YI
        # begin debug
        # plt.plot(subX, ecdf1YI)
        # plt.plot(subX, ecdf2YI)
        # plt.show()
        # pdb.set_trace()
        # end debug
        return torch.tensor(subX), torch.tensor(subY)

    t0 = cifTimes.min()
    tf = cifTimes.max()
    dt = cifTimes[1]-cifTimes[0]
    T = tf-t0
    print("Processing given ISIs")
    zExp = timeRescalingForUnbinnedSpikes(spikesTimes=spikesTimes, cifValues=cifValues, t0=t0, tf=tf, dt=dt)
    zSim = None
    sampler = Sampler()
    for i in range(gamma):
        print("Processing iter {:d}/{:d}".format(i, gamma-1))
        simSpikesTimes = sampler.sampleInhomogeneousPP_timeRescaling(cifTimes=cifTimes, cifValues=cifValues, T=T)
        zSimIter = timeRescalingForUnbinnedSpikes(spikesTimes=simSpikesTimes, cifValues=cifValues, t0=t0, tf=tf, dt=dt)
        if zSim is None:
            zSim = zSimIter
        else:
            zSim = torch.cat((zSim, zSimIter))
    expECDFx, _ = torch.sort(zExp)
    expECDFy = torch.linspace(0, 1, len(zExp))
    simECDFx, _ = torch.sort(zSim)
    simECDFy = torch.linspace(0, 1, len(zSim))
    diffECDFsX, diffECDFsY = subtractECDFs(ecdf1X=expECDFx, ecdf1Y=expECDFy, ecdf2X=simECDFx, ecdf2Y=simECDFy)

    Nexp = len(zExp)
    Nsim = len(zSim)
    cb = 1.36*math.sqrt((Nexp+Nsim)/(Nexp*Nsim))

    return diffECDFsX, diffECDFsY, expECDFx, expECDFy, simECDFx, simECDFy, cb

def timeRescalingAnalyticalCorrection(Y, pk, eps=1e-10):
    indicesMoreThanOneSpikes = (Y>1).nonzero()
    if len(indicesMoreThanOneSpikes)>0:
        warnings.warn("Found more than one spike in {:d} bins".format(len(indicesMoreThanOneSpikes)))
        Y[indicesMoreThanOneSpikes] = 1.0
    pk = torch.max(pk, torch.tensor([eps], dtype=torch.double))
    qk = -torch.log(1-pk)
    # make the rescaled times
    n = int(torch.sum(Y).item())
    rISIs = torch.zeros(n-1)
    spikeindicies = Y.nonzero()
    for r in range(n-1):
        ind1 = spikeindicies[r].item()
        ind2 = spikeindicies[r+1].item()
        aSum = torch.sum(qk[(ind1+1):ind2]).item()
        # delta = -(1.0/qk[ind2].item())*math.log(1-random.random()*(1-math.exp(-qk[ind2].item())))
        # total = total + qk[ind2].item()*delta
        delta = -math.log(1-random.random()*(1-math.exp(-qk[ind2].item())))
        rISIs[r] = aSum + delta
    return rISIs

def KSTestTimeRescalingAnalyticalCorrection(Y, pk, eps=1e-10):
    '''
    Y: binary sequence for the discrete time point process
    pk: event probability in each time been based on a conditional cif function

    xks: uniform CDF
    rst: time rescaled times
    cb: 95% confidence bounds
    n: number of spikes

    KS-test based on time-rescalling theorem =====================
    Use DT Correction for Time Rescaling Theorem - <NAME> (2010)
    discrete time conditional probabilities "pk"  where 0<=pk<= 1
    '''
    rISIs = timeRescalingAnalyticalCorrection(Y=Y, pk=pk, eps=eps) # rescaled inter-spike intervals
    utRISIs = 1-torch.exp(-rISIs)              # uniform-transformed srescaled inter-spike intervals
    srISIs, _ = torch.sort(rISIs)                 # sorted rescaled inter-spike intervals
    utSRISIs = 1-torch.exp(-srISIs)              # uniform-transformed sorted rescaled inter-spike intervals
    n = len(rISIs)
    # dt = 1/(n-1)
    # xks = torch.linspace(0.5*dt, 1-0.5*dt, 1.0/dt)
    k = torch.linspace(start=1, end=n, steps=n)
    uCDF = (k-0.5)/n                            # uniform cummulative distribution function
    cb = 1.36/math.sqrt(n)                      # 95% confidence bounds for large sample sizes (see Brown et al., 2001)
    return utSRISIs, uCDF, cb, utRISIs

def KSTestTimeRescalingAnalyticalCorrectionUnbinned(spikesTimes, cifValues, t0, tf, dt, eps=1e-10):
        pk = cifValues*dt
        bins = torch.arange(t0-dt/2, tf+dt/2, dt)
        # start binning spikes using pandas
        cutRes, _ = pd.cut(spikesTimes, bins=bins, retbins=True)
        Y = torch.from_numpy(cutRes.value_counts().values)
        # end binning spikes using pandas
        utSRISIs, uCDF, cb, utRISIs = KSTestTimeRescalingAnalyticalCorrection(Y=Y, pk=pk, eps=eps)
        return utSRISIs, uCDF, cb, utRISIs


import pdb
import math
import warnings
import random
import pandas as pd
import numpy as np
import torch
import matplotlib.pyplot as plt
from .sampler import Sampler

def timeRescaling(Y, pk):
    n = int(torch.sum(Y).item())
    rISIs = torch.zeros(n-1, dtype=torch.double)
    spikeindicies = Y.nonzero()
    for r in range(n-1):
        ind1 = spikeindicies[r].item()
        ind2 = spikeindicies[r+1].item()
        rISIs[r] = torch.sum(pk[ind1+1:ind2+1]).item()
    return rISIs

def timeRescalingForUnbinnedSpikes(spikesTimes, cifValues, t0, tf, dt):
    pk = cifValues*dt
    bins = torch.arange(t0-dt/2, tf+dt/2, dt)
    # start binning spikes using pandas
    cutRes, _ = pd.cut(spikesTimes, bins=bins, retbins=True)
    Y = torch.from_numpy(cutRes.value_counts().values)
    # end binning spikes using pandas
    indicesMoreThanOneSpikes = (Y>1).nonzero()
    if len(indicesMoreThanOneSpikes)>0:
        warnings.warn("Found more than one spike in {:d} bins".format(len(indicesMoreThanOneSpikes)))
        Y[indicesMoreThanOneSpikes] = 1.0
    zExp = timeRescaling(Y=Y, pk=pk)
    return zExp

def KSTestTimeRescalingNumericalCorrection(spikesTimes, cifTimes, cifValues, gamma):
    # from Haslinger et al., 2010, p. 2492,
    # Procedure for Numerical Correction

    def subtractECDFs(ecdf1X, ecdf1Y, ecdf2X, ecdf2Y):
        subX = torch.sort(torch.cat((ecdf1X, ecdf2X)))[0].numpy()
        ecdf1YI = np.interp(subX, ecdf1X.numpy(), ecdf1Y.numpy())
        ecdf2YI = np.interp(subX, ecdf2X.numpy(), ecdf2Y.numpy())
        subY = ecdf1YI-ecdf2YI
        # begin debug
        # plt.plot(subX, ecdf1YI)
        # plt.plot(subX, ecdf2YI)
        # plt.show()
        # pdb.set_trace()
        # end debug
        return torch.tensor(subX), torch.tensor(subY)

    t0 = cifTimes.min()
    tf = cifTimes.max()
    dt = cifTimes[1]-cifTimes[0]
    T = tf-t0
    print("Processing given ISIs")
    zExp = timeRescalingForUnbinnedSpikes(spikesTimes=spikesTimes, cifValues=cifValues, t0=t0, tf=tf, dt=dt)
    zSim = None
    sampler = Sampler()
    for i in range(gamma):
        print("Processing iter {:d}/{:d}".format(i, gamma-1))
        simSpikesTimes = sampler.sampleInhomogeneousPP_timeRescaling(cifTimes=cifTimes, cifValues=cifValues, T=T)
        zSimIter = timeRescalingForUnbinnedSpikes(spikesTimes=simSpikesTimes, cifValues=cifValues, t0=t0, tf=tf, dt=dt)
        if zSim is None:
            zSim = zSimIter
        else:
            zSim = torch.cat((zSim, zSimIter))
    expECDFx, _ = torch.sort(zExp)
    expECDFy = torch.linspace(0, 1, len(zExp))
    simECDFx, _ = torch.sort(zSim)
    simECDFy = torch.linspace(0, 1, len(zSim))
    diffECDFsX, diffECDFsY = subtractECDFs(ecdf1X=expECDFx, ecdf1Y=expECDFy, ecdf2X=simECDFx, ecdf2Y=simECDFy)

    Nexp = len(zExp)
    Nsim = len(zSim)
    cb = 1.36*math.sqrt((Nexp+Nsim)/(Nexp*Nsim))

    return diffECDFsX, diffECDFsY, expECDFx, expECDFy, simECDFx, simECDFy, cb

def timeRescalingAnalyticalCorrection(Y, pk, eps=1e-10):
    indicesMoreThanOneSpikes = (Y>1).nonzero()
    if len(indicesMoreThanOneSpikes)>0:
        warnings.warn("Found more than one spike in {:d} bins".format(len(indicesMoreThanOneSpikes)))
        Y[indicesMoreThanOneSpikes] = 1.0
    pk = torch.max(pk, torch.tensor([eps], dtype=torch.double))
    qk = -torch.log(1-pk)
    # make the rescaled times
    n = int(torch.sum(Y).item())
    rISIs = torch.zeros(n-1)
    spikeindicies = Y.nonzero()
    for r in range(n-1):
        ind1 = spikeindicies[r].item()
        ind2 = spikeindicies[r+1].item()
        aSum = torch.sum(qk[(ind1+1):ind2]).item()
        # delta = -(1.0/qk[ind2].item())*math.log(1-random.random()*(1-math.exp(-qk[ind2].item())))
        # total = total + qk[ind2].item()*delta
        delta = -math.log(1-random.random()*(1-math.exp(-qk[ind2].item())))
        rISIs[r] = aSum + delta
    return rISIs

def KSTestTimeRescalingAnalyticalCorrection(Y, pk, eps=1e-10):
    '''
    Y: binary sequence for the discrete time point process
    pk: event probability in each time been based on a conditional cif function

    xks: uniform CDF
    rst: time rescaled times
    cb: 95% confidence bounds
    n: number of spikes

    KS-test based on time-rescalling theorem =====================
    Use DT Correction for Time Rescaling Theorem - <NAME> (2010)
    discrete time conditional probabilities "pk"  where 0<=pk<= 1
    '''
    rISIs = timeRescalingAnalyticalCorrection(Y=Y, pk=pk, eps=eps) # rescaled inter-spike intervals
    utRISIs = 1-torch.exp(-rISIs)              # uniform-transformed srescaled inter-spike intervals
    srISIs, _ = torch.sort(rISIs)                 # sorted rescaled inter-spike intervals
    utSRISIs = 1-torch.exp(-srISIs)              # uniform-transformed sorted rescaled inter-spike intervals
    n = len(rISIs)
    # dt = 1/(n-1)
    # xks = torch.linspace(0.5*dt, 1-0.5*dt, 1.0/dt)
    k = torch.linspace(start=1, end=n, steps=n)
    uCDF = (k-0.5)/n                            # uniform cummulative distribution function
    cb = 1.36/math.sqrt(n)                      # 95% confidence bounds for large sample sizes (see Brown et al., 2001)
    return utSRISIs, uCDF, cb, utRISIs

def KSTestTimeRescalingAnalyticalCorrectionUnbinned(spikesTimes, cifValues, t0, tf, dt, eps=1e-10):
        pk = cifValues*dt
        bins = torch.arange(t0-dt/2, tf+dt/2, dt)
        # start binning spikes using pandas
        cutRes, _ = pd.cut(spikesTimes, bins=bins, retbins=True)
        Y = torch.from_numpy(cutRes.value_counts().values)
        # end binning spikes using pandas
        utSRISIs, uCDF, cb, utRISIs = KSTestTimeRescalingAnalyticalCorrection(Y=Y, pk=pk, eps=eps)
        return utSRISIs, uCDF, cb, utRISIs

# start binning spikes using pandas # end binning spikes using pandas # from Haslinger et al., 2010, p. 2492, # Procedure for Numerical Correction # begin debug # plt.plot(subX, ecdf1YI) # plt.plot(subX, ecdf2YI) # plt.show() # pdb.set_trace() # end debug # make the rescaled times # delta = -(1.0/qk[ind2].item())*math.log(1-random.random()*(1-math.exp(-qk[ind2].item()))) # total = total + qk[ind2].item()*delta Y: binary sequence for the discrete time point process     pk: event probability in each time been based on a conditional cif function      xks: uniform CDF     rst: time rescaled times     cb: 95% confidence bounds     n: number of spikes      KS-test based on time-rescalling theorem =====================     Use DT Correction for Time Rescaling Theorem - <NAME> (2010)     discrete time conditional probabilities "pk"  where 0<=pk<= 1 # rescaled inter-spike intervals # uniform-transformed srescaled inter-spike intervals # sorted rescaled inter-spike intervals # uniform-transformed sorted rescaled inter-spike intervals # dt = 1/(n-1) # xks = torch.linspace(0.5*dt, 1-0.5*dt, 1.0/dt) # uniform cummulative distribution function # 95% confidence bounds for large sample sizes (see Brown et al., 2001) # start binning spikes using pandas # end binning spikes using pandas