"""

    This script was used to produce a random first initial layer on which we compare the simulation results and the fit.
        Moreover: the value of inhibitors and activators are selected as random from a law with two gaussian peak, one around 10**(-4) and one around 10**(-8)

    After the simulation, we randomly select an output neuron.
    On this neuron we display the same diagram where activator (inhibitors) are computed as the sum of activations (inhibitions).

"""
import numpy as np
import os
from simulOfBioNN.parseUtils.parser import generateNeuralNetwork,generateTemplateNeuralNetwork,read_file
from simulOfBioNN.simulNN.simulator import executeODESimulation
from simulOfBioNN.odeUtils.systemEquation import fPythonSparse
from simulOfBioNN.odeUtils.utils import readAttribute,obtainTemplateArray,obtainOutputArray
from simulOfBioNN.plotUtils.adaptivePlotUtils import colorDiagram,neuronPlot,plotEvolution,fitComparePlot
from simulOfBioNN.smallNetworkSimul.templateModelRandomNetworkSimul import computeCP as oldComputCP
from simulOfBioNN.smallNetworkSimul.templateModelRandomNetworkSimul import computeEquilibriumValue as oldComputeEquilibriumValue
from simulOfBioNN.smallNetworkSimul.templateModelRandomNetworkSimul import func as oldAllEquilibriumFunc
from scipy.optimize import minimize,root,brentq
import sys
import time
import pandas

def obtainBornSup(k6,kdT,kdI,Kactiv0,Kinhib0,Cactiv0,Cinhib0,E0,X0,masks):
    """
        Given approximate only for the enzyme competition term (cp), we compute the next approximate using G.
        The real value for the competitions term verify cp = G(cp,initialConditions)
    :param cp: float, competition over tempalte.
    :param E0: float, initial concentration in the enzyme
    :param X0: nbrInputs array, contains initial value for the inputs
    :param masks: masks giving the network topology.
    :return:
    """
    max_cp = 1
    olderX = [np.zeros(m.shape[1]) for m in masks]
    for layeridx,layer in enumerate(masks):
        layerEq = np.zeros(layer.shape[1])
        if(layeridx==0):
            for inpIdx in range(layer.shape[1]):
                #compute of Kactivs,Kinhibs;
                Kactivs = np.where(layer[:,inpIdx]>0,Kactiv0[layeridx][:,inpIdx],0) #This is also a matrix element wise multiplication
                Kinhibs = np.where(layer[:,inpIdx]<0,Kinhib0[layeridx][:,inpIdx],0)
                #compute of "weights": sum of kactivs and kinhibs
                w_inpIdx = np.sum(Kactivs)+np.sum(Kinhibs)
                max_cp += w_inpIdx*X0[inpIdx]
                # saving values
                layerEq[inpIdx] = X0[inpIdx]
            olderX[layeridx] = layerEq
        else:
            for inpIdx in range(layer.shape[1]):
                #compute of Cactivs,Cinhibs, the denominator marks the template's variation from equilibrium
                #Terms for the previous layers
                CactivsOld = np.where(masks[layeridx-1][inpIdx,:]>0,Cactiv0[layeridx-1][inpIdx],0)
                CinhibsOld = np.where(masks[layeridx-1][inpIdx,:]<0,Cinhib0[layeridx-1][inpIdx],0)
                Inhib = np.sum(CinhibsOld*olderX[layeridx-1]/kdT[layeridx-1][inpIdx])
                #computing of new equilibrium
                x_eq = np.sum(CactivsOld*olderX[layeridx-1]/kdT[layeridx-1][inpIdx])
                layerEq[inpIdx] = x_eq
                #compute of Kactivs,Kinhibs, for the current layer:
                Kactivs = np.where(layer[:,inpIdx]>0,Kactiv0[layeridx][:,inpIdx],0) #This is also a matrix element wise multiplication
                Kinhibs = np.where(layer[:,inpIdx]<0,Kinhib0[layeridx][:,inpIdx],0)
                #Adding, to the competition over enzyme, the complex formed in this layer by this input.
                firstComplex = np.sum(np.where(layer[:,inpIdx]>0,Kactivs*x_eq,np.where(layer[:,inpIdx]<0,Kinhibs*x_eq,0)))
                #We must also add the effect of pseudoTempalte enzymatic complex in the previous layers which can't be computed previously because we missed x_eq
                Inhib2 = np.sum(CinhibsOld*olderX[layeridx-1]/(kdT[layeridx-1][inpIdx]*k6[layeridx-1][inpIdx]))
                max_cp += firstComplex + Inhib2/E0*x_eq
            olderX[layeridx] = layerEq
    #Finally we must add the effect of pseudoTemplate enzymatic complex in the last layers
    for outputsIdx in range(masks[-1].shape[0]):
        Cinhibs = np.where(masks[-1][outputsIdx,:]<0,Cinhib0[-1][outputsIdx],0)
        Cactivs = np.where(masks[-1][outputsIdx,:]>0,Cactiv0[-1][outputsIdx],0)
        x_eq = np.sum(Cactivs*olderX[-1]/(kdI[-1][outputsIdx]))
        Inhib2 = np.sum(Cinhibs*olderX[-1]/(kdT[-1][outputsIdx]*k6[-1][outputsIdx]))
        max_cp += Inhib2/E0*x_eq
    return max_cp

def cpEquilibriumFunc(cp,k6,kdT,kdI,Kactiv0,Kinhib0,Cactiv0,Cinhib0,E0,X0,masks):
    """
        Given approximate only for the enzyme competition term (cp), we compute the next approximate using G.
        The real value for the competitions term verify cp = G(cp,initialConditions)
    :param cp: float, competition over tempalte.
    :param E0: float, initial concentration in the enzyme
    :param X0: nbrInputs array, contains initial value for the inputs
    :param masks: masks giving the network topology.
    :return:
    """
    new_cp = 1
    olderX = [np.zeros(m.shape[1]) for m in masks]
    for layeridx,layer in enumerate(masks):
        layerEq = np.zeros(layer.shape[1])
        if(layeridx==0):
            for inpIdx in range(layer.shape[1]):
                #compute of Kactivs,Kinhibs;
                Kactivs = np.where(layer[:,inpIdx]>0,Kactiv0[layeridx][:,inpIdx],0) #This is also a matrix element wise multiplication
                Kinhibs = np.where(layer[:,inpIdx]<0,Kinhib0[layeridx][:,inpIdx],0)
                #compute of "weights": sum of kactivs and kinhibs
                w_inpIdx = np.sum(Kactivs)+np.sum(Kinhibs)
                x_eq = X0[inpIdx]/(1+E0*w_inpIdx/cp)
                # update for fixed point:
                new_cp += w_inpIdx*x_eq
                # saving values
                layerEq[inpIdx] = x_eq
            olderX[layeridx] = layerEq
        else:
            for inpIdx in range(layer.shape[1]):

                #compute of Cactivs,Cinhibs, the denominator marks the template's variation from equilibrium
                #Terms for the previous layers
                CactivsOld = np.where(masks[layeridx-1][inpIdx,:]>0,Cactiv0[layeridx-1][inpIdx],0)
                CinhibsOld = np.where(masks[layeridx-1][inpIdx,:]<0,Cinhib0[layeridx-1][inpIdx],0)
                Inhib = np.sum(CinhibsOld*olderX[layeridx-1]/kdT[layeridx-1][inpIdx])
                #computing of new equilibrium
                x_eq = np.sum(CactivsOld*olderX[layeridx-1]/(kdI[layeridx-1][inpIdx]*cp+Inhib/cp))
                layerEq[inpIdx] = x_eq
                #compute of Kactivs,Kinhibs, for the current layer:
                Kactivs = np.where(layer[:,inpIdx]>0,Kactiv0[layeridx][:,inpIdx],0) #This is also a matrix element wise multiplication
                Kinhibs = np.where(layer[:,inpIdx]<0,Kinhib0[layeridx][:,inpIdx],0)
                #Adding, to the competition over enzyme, the complex formed in this layer by this input.
                firstComplex = np.sum(np.where(layer[:,inpIdx]>0,Kactivs*x_eq,np.where(layer[:,inpIdx]<0,Kinhibs*x_eq,0)))
                #We must also add the effect of pseudoTempalte enzymatic complex in the previous layers which can't be computed previously because we missed x_eq
                Inhib2 = np.sum(CinhibsOld*olderX[layeridx-1]/(kdT[layeridx-1][inpIdx]*k6[layeridx-1][inpIdx]))
                new_cp += firstComplex + Inhib2/(E0*cp)*x_eq
            olderX[layeridx] = layerEq
    #Finally we must add the effect of pseudoTemplate enzymatic complex in the last layers
    for outputsIdx in range(masks[-1].shape[0]):
        Cinhibs = np.where(masks[-1][outputsIdx,:]<0,Cinhib0[-1][outputsIdx],0)
        Cactivs = np.where(masks[-1][outputsIdx,:]>0,Cactiv0[-1][outputsIdx],0)

        Inhib = np.sum(Cinhibs*olderX[-1]/kdT[-1][outputsIdx])
        x_eq = np.sum(Cactivs*olderX[-1]/(kdI[-1][outputsIdx]*cp+Inhib/cp))
        Inhib2 = np.sum(Cinhibs*olderX[-1]/(kdT[-1][outputsIdx]*k6[-1][outputsIdx]))
        new_cp += Inhib2/(E0*cp)*x_eq
    return cp - new_cp

def allEquilibriumFunc(cps,k6,kdT,kdI,Kactiv0,Kinhib0,Cactiv0,Cinhib0,E0,X0,masks):
    """
        Given approximate for the different competitions term (in the vector cps), we compute the next approximate using G.
        The real value for the competitions term verify cps = G(cps,initialConditions)
    :param cps: 1+nbrTemplate array. In the first one we store the competition on the enzyme.
                                     In the other the competition on the template species.
    :param E0: float, initial concentration in the enzyme
    :param X0: nbrInputs array, contains initial value for the inputs
    :param masks: masks giving the network topology.
    :return:
    """
    # We move from an array to a list of 2d-array for the competition over each template:

    cp=cps[0]
    cpt = [np.reshape(cps[(l*m.shape[0]*m.shape[1]+1):((l+1)*m.shape[0]*m.shape[1]+1)],(m.shape[0],m.shape[1])) for l,m in enumerate(masks)]
    new_cpt = [np.zeros(l.shape)+1 for l in cpt]
    new_cp = 1

    olderX = [np.zeros(m.shape[1]) for m in masks]
    for layeridx,layer in enumerate(masks):
        layerEq = np.zeros(layer.shape[1])
        if(layeridx==0):
            for inpIdx in range(layer.shape[1]):
                #compute of Kactivs,Kinhibs;
                Kactivs = np.where(layer[:,inpIdx]>0,Kactiv0[layeridx][:,inpIdx]/cpt[layeridx][:,inpIdx],0) #This is also a matrix element wise multiplication
                Kinhibs = np.where(layer[:,inpIdx]<0,Kinhib0[layeridx][:,inpIdx]/cpt[layeridx][:,inpIdx],0)
                #compute of "weights": sum of kactivs and kinhibs
                w_inpIdx = np.sum(Kactivs)+np.sum(Kinhibs)
                x_eq = X0[inpIdx]/(1+E0*w_inpIdx/cp)
                # update for fixed point:
                new_cp += w_inpIdx*x_eq
                # saving values
                layerEq[inpIdx] = x_eq
            new_cpt[layeridx] = np.where(layer>0,(1+k1M[layeridx]*E0*layerEq/cp),np.where(layer<0,(1+k3M[layeridx]*E0*layerEq/cp),1))
            olderX[layeridx] = layerEq
        else:
            for inpIdx in range(layer.shape[1]):

                #compute of Cactivs,Cinhibs, the denominator marks the template's variation from equilibrium
                #Terms for the previous layers
                CactivsOld = np.where(masks[layeridx-1][inpIdx,:]>0,Cactiv0[layeridx-1][inpIdx]/cpt[layeridx-1][inpIdx],0)
                CinhibsOld = np.where(masks[layeridx-1][inpIdx,:]<0,Cinhib0[layeridx-1][inpIdx]/cpt[layeridx-1][inpIdx],0)
                Inhib = np.sum(CinhibsOld*olderX[layeridx-1]/kdT[layeridx-1][inpIdx])
                #computing of new equilibrium
                x_eq = np.sum(CactivsOld*olderX[layeridx-1]/(kdI[layeridx-1][inpIdx]*cp+Inhib/cp))
                layerEq[inpIdx] = x_eq

                #compute of Kactivs,Kinhibs, for the current layer:
                Kactivs = np.where(layer[:,inpIdx]>0,Kactiv0[layeridx][:,inpIdx]/cpt[layeridx][:,inpIdx],0) #This is also a matrix element wise multiplication
                Kinhibs = np.where(layer[:,inpIdx]<0,Kinhib0[layeridx][:,inpIdx]/cpt[layeridx][:,inpIdx],0)
                #Adding, to the competition over enzyme, the complex formed in this layer by this input.
                firstComplex = np.sum(np.where(layer[:,inpIdx]>0,Kactivs*x_eq,np.where(layer[:,inpIdx]<0,Kinhibs*x_eq,0)))
                #We must also add the effect of pseudoTempalte enzymatic complex in the previous layers which can't be computed previously because we missed x_eq
                Inhib2 = np.sum(CinhibsOld*olderX[layeridx-1]/(kdT[layeridx-1][inpIdx]*k6[layeridx-1][inpIdx]))
                new_cp += firstComplex + Inhib2/(E0*cp)*x_eq
            new_cpt[layeridx] = np.where(layer>0,(1+k1M[layeridx]*E0*layerEq/cp),np.where(layer<0,(1+k3M[layeridx]*E0*layerEq/cp),1))
            olderX[layeridx] = layerEq
    #Finally we must add the effect of pseudoTemplate enzymatic complex in the last layers
    for outputsIdx in range(masks[-1].shape[0]):
        Cinhibs = np.where(masks[-1][outputsIdx,:]<0,Cinhib0[-1][outputsIdx]/cpt[-1][outputsIdx],0)
        Cactivs = np.where(masks[-1][outputsIdx,:]>0,Cactiv0[-1][outputsIdx]/cpt[-1][outputsIdx],0)

        Inhib = np.sum(Cinhibs*olderX[-1]/kdT[-1][outputsIdx])
        x_eq = np.sum(Cactivs*olderX[-1]/(kdI[-1][outputsIdx]*cp+Inhib/cp))
        Inhib2 = np.sum(Cinhibs*olderX[-1]/(kdT[-1][outputsIdx]*k6[-1][outputsIdx]))
        new_cp += Inhib2/(E0*cp)*x_eq

    #we know that the root should be larger than 1 so we artificially add a root barrier at 1.
    diff_cps = np.zeros(cps.shape)
    if new_cp>=1:
        diff_cps[0] = cp - new_cp
    else:
        diff_cps[0] = 10**3
    for l,m in enumerate(masks):
        layer_new_cpt = np.reshape(new_cpt[l],(m.shape[0]*m.shape[1]))
        layer_old_cpt = cps[(l*m.shape[0]*m.shape[1]+1):((l+1)*m.shape[0]*m.shape[1]+1)]
        diff_cps[(l*m.shape[0]*m.shape[1]+1):((l+1)*m.shape[0]*m.shape[1]+1)] = layer_old_cpt - np.where(layer_new_cpt<1,layer_old_cpt + 10**3, layer_new_cpt)

    return diff_cps

def for_scalar(cps,k6,kdT,kdI,Kactiv0,Kinhib0,Cactiv0,Cinhib0,E0,X0,masks):
    return np.sum(np.abs(allEquilibriumFunc(cps,k6,kdT,kdI,Kactiv0,Kinhib0,Cactiv0,Cinhib0,E0,X0,masks)))

def computeCPonly(k1,k1n,k2,k3,k3n,k4,k5,k5n,k6,kdI,kdT,TA0,TI0,E0,X0,masks,fittedValue=None):
    """
         This function computes the competition's value by solving a fixed point equation.
         It is based on the most simplest chemical model for the template model: no endonuclease; polymerase and nickase are considered together.
     :param k1, and the others k are the reactions constants
     :param X: array with initial concentration of the first layer.
     :param masks:
     :param fittedValue: value obtained with an analytical model that is an upper bound on the real value. if not provided we use 10**6
     :return:
     """
    # Computation of the different constants that are required.
    k1M = [k1[l]/(k1n[l]+k2[l]) for l in range(len(k1))] #(matrix operations)
    k3M = [k3[l]/(k3n[l]+k4[l]) for l in range(len(k3))]
    k5M = [k5[l]/(k5n[l]+k6[l]) for l in range(len(k5))]

    Kactiv0 = [k1M[l]*TA0[l] for l in range(len(k1M))] # element to element product
    Kinhib0 = [k3M[l]*TI0[l] for l in range(len(k3M))]
    Cactiv0 = [k2[l]*k1M[l]*TA0[l]*E0 for l in range(len(k1M))]
    Cinhib0 =[k6[l]*k5M[l]*k4[l]*k3M[l]*TI0[l]*E0*E0 for l in range(len(k3M))]

    if fittedValue is None:
        cp0max=obtainBornSup(k6,kdT,kdI,Kactiv0,Kinhib0,Cactiv0,Cinhib0,E0,X0,masks)# we start from an upper bound
    else:
        cp0max=fittedValue
    t0=time.time()
    computedCp = brentq(cpEquilibriumFunc,1,cp0max,args=(k6,kdT,kdI,Kactiv0,Kinhib0,Cactiv0,Cinhib0,E0,X0,masks))
    print("Ended brentq methods in "+str(time.time()-t0))
    return computedCp

def computeCP(k1,k1n,k2,k3,k3n,k4,k5,k5n,k6,kdT,kdI,TA0,TI0,E0,X0,masks,initValue=None):
    """
        This function computes the competition's value by solving a fixed point equation.
        It is based on the most simplest chemical model for the template model: no endonuclease; polymerase and nickase are considered together.
    :param k1, and the others k are the reactions constants
    :param X0: array with initial concentration of the first layer.
    :param masks:
    :param initValue: initial value for cp, the competition over the enzymes. Can be computed with brentq method by fixing the competition over template...
    :return:
    """
    # Computation of the different constants that are required.
    k1M = [k1[l]/(k1n[l]+k2[l]) for l in range(len(k1))] #(matrix operations)
    k3M = [k3[l]/(k3n[l]+k4[l]) for l in range(len(k3))]
    k5M = [k5[l]/(k5n[l]+k6[l]) for l in range(len(k5))]

    Kactiv0 = [k1M[l]*TA0[l] for l in range(len(k1M))] # element to element product
    Kinhib0 = [k3M[l]*TI0[l] for l in range(len(k3M))]
    Cactiv0 = [k2[l]*k1M[l]*TA0[l]*E0 for l in range(len(k1M))]
    Cinhib0 =[k6[l]*k5M[l]*k4[l]*k3M[l]*TI0[l]*E0*E0 for l in range(len(k3M))]

    if initValue is None:
        cps0min=np.concatenate(([1],np.zeros(np.sum([m.shape[0]*m.shape[1] for m in masks]))+1)) # we start from an inferior bound
    else:
        cps0min=np.concatenate(([initValue],np.zeros(np.sum([m.shape[0]*m.shape[1] for m in masks]))+1))
    t0=time.time()
    #computedCp = minimize(for_scalar,cps0min,args=(k2,k4,k6,kdT,kdI,Kactiv0,Kinhib0,Cactiv0,Cinhib0,TA0,TI0,E0,X0,masks),method="L-BFGS-B",bounds=[(1,None) for _ in cps0min])
    computedCp = root(allEquilibriumFunc,cps0min,args=(k6,kdT,kdI,Kactiv0,Kinhib0,Cactiv0,Cinhib0,E0,X0,masks))
    print("ended root methods in "+str(time.time()-t0))
    if not computedCp["success"]:
        print(computedCp["message"])
    return computedCp["x"]

def computeEquilibriumValue(cps,k1,k1n,k2,k3,k3n,k4,k5,k5n,k6,kdT,kdI,TA0,TI0,E0,X0,masks,observed=None):
    """
        Given cp, compute the equilibrium values for all nodes in solutions.
    :param cps: Value for the competitions over enzyme and template obtained by fixed point strategy.
    :param observed: tuple, default to None. If provided, we only give the value for the species at the position observed.
    :return:
    """
    k1M = [k1[l]/(k1n[l]+k2[l]) for l in range(len(k1))] #(matrix operations)
    k3M = [k3[l]/(k3n[l]+k4[l]) for l in range(len(k3))]
    k5M = [k5[l]/(k5n[l]+k6[l]) for l in range(len(k5))]
    Kactiv0 = [k1M[l]*TA0[l] for l in range(len(k1M))] # element to element product
    Kinhib0 = [k3M[l]*TI0[l] for l in range(len(k3M))]
    Cactiv0 = [k2[l]*k1M[l]*TA0[l]*E0 for l in range(len(k1M))]
    Cinhib0 =[k6[l]*k5M[l]*k4[l]*k3M[l]*TI0[l]*E0*E0 for l in range(len(k3M))]

    # We move from an array to a list of 2d-array for the competition over each template:
    cp=cps[0]
    cpt = [np.reshape(cps[(l+1):(l+1+m.shape[0]*m.shape[1])],(m.shape[0],m.shape[1])) for l,m in enumerate(masks)]
    equilibriumValue = [np.zeros(m.shape[1]) for m in masks]+[np.zeros(masks[-1].shape[0])]
    for layeridx,layer in enumerate(masks):
        layerEq = np.zeros(layer.shape[1])
        if(layeridx==0):
            for inpIdx in range(layer.shape[1]):
                #compute of Kactivs,Kinhibs;
                Kactivs = np.where(layer[:,inpIdx]>0,Kactiv0[layeridx][:,inpIdx]/cpt[layeridx][:,inpIdx],0) #This is also a matrix element wise multiplication
                Kinhibs = np.where(layer[:,inpIdx]<0,Kinhib0[layeridx][:,inpIdx]/cpt[layeridx][:,inpIdx],0)
                #compute of "weights": sum of kactivs and kinhibs
                w_inpIdx = np.sum(Kactivs)+np.sum(Kinhibs)
                x_eq = X0[inpIdx]/(1+E0*w_inpIdx/cp)
                # saving values
                layerEq[inpIdx] = x_eq
            equilibriumValue[layeridx] = layerEq
        else:
            for inpIdx in range(layer.shape[1]):
                #compute of Cactivs,Cinhibs, the denominator marks the template's variation from equilibrium
                #Terms for the previous layers
                CactivsOld = np.where(masks[layeridx-1][inpIdx,:]>0,Cactiv0[layeridx-1][inpIdx]/cpt[layeridx-1][inpIdx],0)
                CinhibsOld = np.where(masks[layeridx-1][inpIdx,:]<0,Cinhib0[layeridx-1][inpIdx]/cpt[layeridx-1][inpIdx],0)

                #compute of Kactivs,Kinhibs, for the current layer:
                Kactivs = np.where(layer[:,inpIdx]>0,Kactiv0[layeridx][:,inpIdx]/cpt[layeridx][:,inpIdx],0) #This is also a matrix element wise multiplication
                Kinhibs = np.where(layer[:,inpIdx]<0,Kinhib0[layeridx][:,inpIdx]/cpt[layeridx][:,inpIdx],0)
                w_nextLayer = E0*(np.sum(Kactivs*k2[layeridx][:,inpIdx])+np.sum(Kinhibs*k4[layeridx][:,inpIdx])) #the input influences the next layer

                Inhib = np.sum(CinhibsOld*equilibriumValue[layeridx-1]/kdT[layeridx-1][inpIdx])
                #computing of new equilibrium
                x_eq = np.sum(CactivsOld*equilibriumValue[layeridx-1]/(kdI[layeridx-1][inpIdx]*cp+Inhib/cp))
                layerEq[inpIdx] = x_eq
            equilibriumValue[layeridx] = layerEq

    #Compute equilibrium for last layers:
    layerEq = np.zeros(masks[-1].shape[0])
    for outputIdx in range(masks[-1].shape[0]):
        CactivsOld = np.where(masks[-1][outputIdx,:]>0,Cactiv0[-1][outputIdx]/cpt[-1][outputIdx],0)
        CinhibsOld = np.where(masks[-1][outputIdx,:]<0,Cinhib0[-1][outputIdx]/cpt[-1][outputIdx],0)

        Inhib = np.sum(CinhibsOld*equilibriumValue[-2]/kdT[-1][outputIdx])
        #computing of new equilibrium
        x_eq = np.sum(CactivsOld*equilibriumValue[-2]/(kdI[-1][outputIdx]*cp+Inhib/cp))
        layerEq[outputIdx]=x_eq
    equilibriumValue[-1] = layerEq

    if observed is not None:
        try:
            assert len(observed)==2
        except:
            raise Exception("please provide a tuple of size two for observed indicating the value to use.")
        return equilibriumValue[observed[0]][observed[1]]
    return equilibriumValue

def _createMask(nbrOutputNodes):
    mask=[]
    for idx in range(nbrOutputNodes):
        zero=np.zeros(2*nbrOutputNodes)
        zero[2*idx]=1
        zero[2*idx+1]=-1
        mask+=[zero]
    return np.array([mask])

def sample_bimodal_distrib(nbrInputs,peak1=-4,peak2=-8,sigma1=1.,sigma2=1.,mixture = 0.5):
    """
        Sample from a bimodal distribution around peak 1 and peak 2 based on gaussian law.
    :param nbrInputs: int, indicate the number of inputs
    :param peak1: int, indicate the power of 10 around which the first peak should be centered
    :param peak2: int, indicate the power of 10 around which the second peak should be centered
    :param sigma1: float, indicate the width in power of 10 around the peak1.
    :param sigma2: float, indicate the width in power of 10 around the peak2.
    :param mixture: proba for the mixture of the two.
    :return:
    """
    mixChoice = np.random.random(nbrInputs)
    gaussian1 = np.random.normal(peak1,sigma1,nbrInputs)
    gaussian2 = np.random.normal(peak2,sigma2,nbrInputs)

    powers = np.where(mixChoice>mixture,gaussian1,gaussian2)
    inputs = [10**p for p in powers]
    return inputs,powers

def _sample_layer_architecture(nbrInputs,nbrOutputs,sparsity=0.5):
    """
        The sparsity indicate the percentage of empty connections.
    :param nbrInputs: number of nodes in the first layer
    :param nbrOutputs: number of nodes in the second layer
    :param sparsity: indicate the percentage of empty connections
    :return:
    """
    choiceEdges = np.random.random((nbrOutputs,nbrInputs)) # choose the number of edges that should be connected.
    choiceActiv = np.random.random((nbrOutputs,nbrInputs)) # choose the edges that are activators (found number below 0.5).
    masks = [np.where(choiceEdges<sparsity,np.zeros(choiceEdges.shape),np.where(choiceActiv>0.5,np.zeros(choiceEdges.shape)+1,np.zeros(choiceEdges.shape)-1))]
    return masks

def _generate_initialConcentration_firstlayer(activatorsOnObserved,inhibitorsOnObserved,masks,nbrInputs,
                                              minLogSpace=-8,maxLogSpace=-4,nbrValue=10):
    """
        Generate the concentrations for the first layer.

    :param minLogSpace: minimal power of 10 for the gaussian mean, default to -8
    :param maxLogSpace: maximal power of 10 for the gaussian mean, default to -4
    :param nbrValue: indicate the number of points that should be taken between [minLogSpace,maxLogspace]
    :return:
    """

    # Generate the first layer concentration:
    #For the activators from the observed node, we sample values in a gaussian disribution, with mean varying on a log scale.
    peaks = np.log(np.logspace( minLogSpace , maxLogSpace , nbrValue))/np.log(10)
    activPowers = [np.random.normal(p,1,len(activatorsOnObserved)) for p in peaks]
    inhibPowers = [np.random.normal(p,1,len(inhibitorsOnObserved)) for p in peaks]
    activInputs = []
    for acp in activPowers:
        activInputs+=[[10**p for p in acp]]
    inhibInputs = []
    for inp in inhibPowers:
        inhibInputs+=[[10**p for p in inp]]
    activInputs = np.array(activInputs)
    inhibInputs = np.array(inhibInputs)

    X1 = np.sum(activInputs,axis=1)
    argsortX1 = np.argsort(X1)
    X1 = X1[argsortX1]
    X2 = np.sum(inhibInputs,axis=1)
    argsortX2 = np.argsort(X2)
    X2 = X2[argsortX2]

    activInputs = activInputs[argsortX1,:]
    inhibInputs = inhibInputs[argsortX2,:]

    # For the remaining inputs we sample from a bimodal distribution:
    myOtherInputs = sample_bimodal_distrib(nbrInputs-len(activatorsOnObserved)-len(inhibitorsOnObserved),sigma1=0.5,sigma2=0.5)[0]
    otherInputs = np.array([myOtherInputs for p in peaks])

    # To compute the competition properly, we must add the concentrations of species for each of the activations or inhibitions
    otherActivInitialC = np.zeros((len(peaks),len(peaks)))
    otherInhibInitialC = np.zeros((len(peaks),len(peaks)))

    x_test=np.zeros((len(peaks),len(peaks),nbrInputs))
    for idxp0,p0 in enumerate(peaks):
        line=np.zeros((len(peaks),nbrInputs))
        for idxp,p in enumerate(peaks):
            for i,idx in enumerate(activatorsOnObserved):
                line[idxp,idx] = activInputs[idxp0,i] #We keep the activation constant for the second axis.
                otherActivInitialC[idxp0,idxp]+= line[idxp, idx] * (np.sum(masks[0][ : , idx] > 0) - 1)
                otherInhibInitialC[idxp0,idxp]+= line[idxp, idx] * np.sum(masks[0][ :, idx] < 0)
            for i,idx in enumerate(inhibitorsOnObserved):
                line[idxp,idx] = inhibInputs[idxp,i]
                otherActivInitialC[idxp0,idxp]+= line[idxp, idx] * np.sum(masks[0][ :, idx] > 0)
                otherInhibInitialC[idxp0,idxp]+= line[idxp, idx] * (np.sum(masks[0][ :, idx] < 0)-1)
            c=0
            for idx in range(nbrInputs):
                if idx not in activatorsOnObserved and idx not in inhibitorsOnObserved:
                    line[idxp,idx] = otherInputs[idxp,c]
                    c=c+1
                    otherActivInitialC[idxp0,idxp]+= line[idxp, idx] * np.sum(masks[0][ :, idx] > 0)
                    otherInhibInitialC[idxp0,idxp]+= line[idxp, idx] * np.sum(masks[0][ :, idx] < 0)
        x_test[idxp0]=line

    x_test = np.reshape(x_test,(len(peaks)*len(peaks),nbrInputs))

    return X1,X2,x_test,otherActivInitialC,otherInhibInitialC

if __name__ == '__main__':

    name = "templateModelTwoLayer/equiNetwork_100_newfit"
    endTime = 10000
    timeStep = 0.1

    # masks=[np.array([[1,-1,0,0],[0,0,1,-1]]),np.array([[1,-1]])]
    # nbrInputs=masks[0].shape[1]
    #
    # activatorsOnObserved = [0]
    # inhibitorsOnObserved = [1]
    # nodeObserved = 0
    nbrInputs=[100,10]
    nbrOutputs=[10,5]
    activatorsOnObserved = []
    inhibitorsOnObserved = []
    while(len(activatorsOnObserved)==0 and len(inhibitorsOnObserved)==0):
        masks=[]
        for idx,nbi in enumerate(nbrInputs):
            masks+= _sample_layer_architecture(nbi,nbrOutputs[idx],sparsity=.5)
        nodeObserved = np.random.randint(0,nbrOutputs[0],1)[0]
        # We look at the activators and inhibitors on the observed node, and verify that they have a length greater than one (while loop).
        for idx,m in enumerate(masks[0][nodeObserved]):
            if m==1:
                activatorsOnObserved += [idx]
            elif m==-1:
                inhibitorsOnObserved += [idx]

    print(masks)
    print("Observed nodes has "+str(len(activatorsOnObserved))+" activators and "+str(len(inhibitorsOnObserved))+" inhibitors")
    print(masks[0][nodeObserved])
    # For debug: we start by working with competition on template from the first node.
    templateObservedIdx = 1 + nodeObserved*masks[0].shape[1] + activatorsOnObserved[0]

    modes = ["verbose","outputEqui"]
    FULL = True
    outputMode = "all"
    outputList = ["X_1_"+str(nodeObserved)] #"all" or None or list of output species
    complexity="simple"
    useEndo = False  # if we want to use the complicated endo model
    useProtectionOnActivator = False
    useEndoOnInputs = False
    useEndoOnOutputs = True
    useDerivativeLeak = True

    leak = 10**(-10)

    layerInit = 10**(-13) #initial concentation value for species in layers
    initValue = 10**(-13) #initial concentration value for all species.
    enzymeInit = 5*10**(-7)
    endoInit = 10**(-5) #only used if useEndo == True
    activInit =  10**(-4)
    inhibInit =  10**(-4)

    if useEndo:
        generateTemplateNeuralNetwork(name,masks,complexity=complexity,useProtectionOnActivator=useProtectionOnActivator,
                                      useEndoOnOutputs=useEndoOnOutputs,useEndoOnInputs=useEndoOnInputs)
    else:
        generateTemplateNeuralNetwork(name,masks,complexity=complexity,endoConstants=None,useProtectionOnActivator=useProtectionOnActivator,
                                      useEndoOnInputs=useEndoOnInputs,useEndoOnOutputs=useEndoOnOutputs)

    X1,X2,x_test,otherActivInitialC,otherInhibInitialC = _generate_initialConcentration_firstlayer(activatorsOnObserved,inhibitorsOnObserved,masks,nbrInputs[0],minLogSpace=-8,maxLogSpace=-4,nbrValue=5)

    # We realised that the rescale factor should be proportionate to the number of edges in order to keep competition low compare to the inhibition.
    # Moreover we observe that rescaling the template rather than the activation is probably better.
    computedRescaleFactor = np.sum([np.sum(m>0)+np.sum(m<0) for m in masks])
    print("Computed rescale factor is "+str(computedRescaleFactor))

    # Finally we observe that rescaling the number of enzyme up can also help to diminish the competition compare to the inhibition as the enzyme appear with a squared power in the last one
    enzymeInit = enzymeInit*(computedRescaleFactor**0.5)
    inhibInit = inhibInit
    activInit = activInit

    initialization_dic={}
    for layer in range(0,len(masks)): ## The first layer need not to be initiliazed
        for node in range(masks[layer].shape[0]):
            initialization_dic["X_"+str(layer+1)+"_"+str(node)] = layerInit
    inhibTemplateNames = obtainTemplateArray(masks=masks,activ=False)
    for k in inhibTemplateNames:
        initialization_dic[k] = inhibInit
    activTemplateNames = obtainTemplateArray(masks=masks,activ=True)
    for k in activTemplateNames:
        initialization_dic[k] = activInit
    initialization_dic["E"] = enzymeInit
    if complexity!="simple":
        initialization_dic["E2"] = enzymeInit
    if complexity!=None and useEndo:
        initialization_dic["Endo"] = endoInit

    if useDerivativeLeak:
        results = executeODESimulation(fPythonSparse, name, x_test, initialization_dic, outputList= outputList,
                                       leak = leak, endTime=endTime, sparse=True, modes=modes,
                                       timeStep=timeStep, initValue= initValue, rescaleFactor=computedRescaleFactor)
    else:
        results = executeODESimulation(fPythonSparse, name, x_test, initialization_dic, outputList= outputList,
                                       leak = 0, endTime=endTime, sparse=True, modes=modes,
                                       timeStep=timeStep, initValue= initValue, rescaleFactor=computedRescaleFactor)

    if("outputPlot" in modes):
        shapeP = len(X1)*len(X2)
        nameDic = results[-1]
        X = results[modes.index("outputPlot")]
        myname=os.path.join(sys.path[0],name)
        if(outputMode=="all"):
            if outputList is None or type(outputList)==str:
                outputList = list(nameDic.keys())
        else:
            outputList = obtainOutputArray(nameDic)
        specialnameDic ={}
        for idx2,e in enumerate(outputList):
            specialnameDic[e] = idx2
        for e in range(X.shape[1]):
            displayX = np.moveaxis(X[:,e,:],0,1)
            plotEvolution(np.arange(0,endTime,timeStep),os.path.join(myname,str(e)), specialnameDic, displayX, wishToDisp=outputList, displaySeparate=True, displayOther=True)

    if("outputEqui" in modes):
        experiment_path = name
        C0=readAttribute(experiment_path,["C0"])["C0"]
        rescaleFactor = readAttribute(experiment_path,["rescaleFactor"])["rescaleFactor"]
        output = results[modes.index("outputEqui")]
        output = np.reshape(output,(len(X1),len(X2)))
        X1 = X1/(C0*rescaleFactor)
        X2 = X2/(C0*rescaleFactor)
        otherActivInitialC = otherActivInitialC/(C0*rescaleFactor)
        otherInhibInitialC = otherInhibInitialC/(C0*rescaleFactor)

        # colorDiagram(X1,X2,output,"Initial concentration of X1","Initial concentration of X2","Equilibrium concentration of the output",figname=os.path.join(experiment_path, "neuralDiagramm.png"),equiPotential=False)
        # neuronPlot(X1,X2,output,figname=os.path.join(experiment_path, "activationLogX1.png"),figname2=os.path.join(experiment_path, "activationLogX2.png"),useLogX = True, doShow= False)
        # neuronPlot(X1,X2,output,figname=os.path.join(experiment_path, "activationX1.png"),figname2=os.path.join(experiment_path, "activationX2.png"),useLogX = False, doShow= False)
        # df = pandas.DataFrame(X1)
        # df.to_csv(os.path.join(experiment_path,"inputX1.csv"))
        # df2 = pandas.DataFrame(X2)
        # df2.to_csv(os.path.join(experiment_path,"inputX2.csv"))

        #we try to fit:
        nbrConstant = int(readAttribute(experiment_path,["Numbers_of_Constants"])["Numbers_of_Constants"])
        if nbrConstant == 12: #only one neuron, it is easy to extract cste values
            k1,k1n,k2,k3,k3n,k4,_,k5,k5n,k6,kd,_=[readAttribute(experiment_path,["k"+str(i)])["k"+str(i)] for i in range(0,nbrConstant)]
        else:
            k1,k1n,k2,k3,k3n,k4,_,k5,k5n,k6,kd,_= [0.9999999999999998,0.1764705882352941,1.0,0.9999999999999998,0.1764705882352941,1.0,0.018823529411764708,0.9999999999999998,0.1764705882352941,1.0,0.018823529411764708,0.018823529411764708]

        TA = activInit/C0
        TI = inhibInit/C0
        E0 = enzymeInit/C0
        k1M = k1/(k1n+k2)
        k5M = k5/(k5n+k6)
        k3M = k3/(k3n+k4)
        Cactiv = k2*k1M*TA*E0
        CInhib = k6*k5M*k4*k3M*TI*E0*E0
        Kactiv = k1M*TA
        Kinhib = k3M*TI
        print("Cactiv value is :"+str(Cactiv))
        print("Cinhib value is :"+str(CInhib))
        print("Kactiv value is :"+str(Kactiv))
        print("Kinhib value is :"+str(Kinhib))
        print("kd is :"+str(kd))
        print("template activator value is:"+str(TA))
        print("template inhibitor value is:"+str(TI))

        #Third fit, using the compute cp:
        myX_test=np.reshape(x_test,(len(X1),len(X2),nbrInputs[0]))/(C0*rescaleFactor)
        competitions = np.zeros((len(X1),len(X2)))
        oldCompetitions = np.zeros((len(X1),len(X2)))
        fitOutput = np.zeros((len(X1), len(X2)))
        oldFitOutput = np.zeros((len(X1), len(X2)))
        styleFit = np.zeros((len(X1),len(X2),1000))
        oldStyleFit = np.zeros((len(X1),len(X2),1000))
        testOfCp = np.logspace(-5,3,1000)/kd
        testOfCpt = np.logspace(-3,-1,1000)/kd
        courbs=[0,int(fitOutput.shape[1]/2),fitOutput.shape[1]-1,int(fitOutput.shape[1]/3),int(2*fitOutput.shape[1]/3)]

        k1 = [np.zeros(m.shape)+k1 for m in masks]
        k1n = [np.zeros(m.shape)+k1n for m in masks]
        k2 = [np.zeros(m.shape)+k2 for m in masks]
        k3 = [np.zeros(m.shape)+k3 for m in masks]
        k3n = [np.zeros(m.shape)+k3n for m in masks]
        k4 = [np.zeros(m.shape)+k4 for m in masks]
        k5 = [np.zeros(m.shape)+k5 for m in masks]
        k5n = [np.zeros(m.shape)+k5n for m in masks]
        k6 = [np.zeros(m.shape)+k6 for m in masks]
        kdT = [np.zeros(m.shape)+kd for m in masks]
        kdI = [np.zeros(m.shape)+kd for m in masks]

        TA0 = [np.where(m>0,TA,0) for m in masks]
        TI0 = [np.where(m<0,TI,0) for m in masks]
        k1M = [k1[l]/(k1n[l]+k2[l]) for l in range(len(k1))] #(matrix operations)
        k3M = [k3[l]/(k3n[l]+k4[l]) for l in range(len(k3))]
        k5M = [k5[l]/(k5n[l]+k6[l]) for l in range(len(k5))]
        Kactiv0 = [k1M[l]*TA0[l] for l in range(len(k1M))] # element to element product
        Kinhib0 = [k3M[l]*TI0[l] for l in range(len(k3M))]
        Cactiv0 = [k2[l]*k1M[l]*TA0[l]*E0 for l in range(len(k1M))]
        Cinhib0 =[k6[l]*k5M[l]*k4[l]*k3M[l]*TI0[l]*E0*E0 for l in range(len(k3M))]

        visuCPT = np.zeros((len(X1),len(X2),1000))
        outCPT = np.zeros((len(X1),len(X2)))
        for idx1,x1 in enumerate(X1):
            for idx2,x2 in enumerate(X2):
                oldCompetitions[idx1,idx2] = oldComputCP(k1[0][0,0], k1n[0][0,0], k2[0][0,0], k3[0][0,0], k3n[0][0,0], k4[0][0,0], k5[0][0,0], k5n[0][0,0], k6[0][0,0], kdT[0][0,0], TA, TI,E0,myX_test[idx1,idx2],masks)
                cpApproxim = computeCPonly(k1, k1n, k2, k3, k3n, k4, k5, k5n, k6, kdT,kdI, TA0, TI0, E0,myX_test[idx1,idx2],masks)
                cps = computeCP(k1, k1n, k2, k3, k3n, k4, k5, k5n, k6, kdT,kdI, TA0, TI0, E0,myX_test[idx1,idx2],masks,initValue=cpApproxim)
                competitions[idx1,idx2] = cps[0]
                outCPT[idx1,idx2] = cps[1]
                oldFitOutput[idx1,idx2] = oldComputeEquilibriumValue(oldCompetitions[idx1,idx2],myX_test[idx1,idx2],masks ,k1[0][0,0], k1n[0][0,0], k2[0][0,0], k3[0][0,0], k3n[0][0,0], k4[0][0,0], k5[0][0,0], k5n[0][0,0], k6[0][0,0], kdT[0][0,0], TA, TI,E0,observed=(1, nodeObserved))
                fitOutput[idx1,idx2] = computeEquilibriumValue(cps,k1,k1n,k2,k3,k3n,k4,k5,k5n,k6,kdT,kdI,TA0,TI0,E0,myX_test[idx1,idx2],masks,observed=(1, nodeObserved))
                if idx2 in courbs and idx1==idx2:
                    testcps = []
                    for x in testOfCp:
                        testcps += [np.concatenate(([x],np.zeros(np.sum([m.shape[0]*m.shape[1] for m in masks]))+1))]
                    styleFit[idx1,idx2] = np.array([allEquilibriumFunc(testofcps,k6,kdT,kdI,Kactiv0,Kinhib0,Cactiv0,Cinhib0,E0,myX_test[idx1,idx2],masks)[0] for testofcps in testcps])
                    oldStyleFit[idx1,idx2] = np.array([oldAllEquilibriumFunc(x*kd,kdT[0][0,0],E0,k2[0][0,0],k4[0][0,0],k6[0][0,0],Cactiv,CInhib,Kactiv,Kinhib,masks,myX_test[idx1,idx2]) for x in testOfCp])
                    testCPTs=[]
                    for x in testOfCpt:
                        cptInitArray = np.zeros(np.sum([m.shape[0]*m.shape[1] for m in masks])-1)+1
                        cptInitArray[templateObservedIdx-1] = x
                        testCPTs += [np.concatenate(([oldCompetitions[idx1,idx2]*kd],cptInitArray))]
                    visuCPT[idx1,idx2] = np.array([allEquilibriumFunc(testCPT,k6,kdT,kdI,Kactiv0,Kinhib0,Cactiv0,Cinhib0,E0,myX_test[idx1,idx2],masks)[templateObservedIdx] for testCPT in testCPTs])
        # comparisons of plot
        fitComparePlot(X1, X2, output, fitOutput, courbs,
                       figname=os.path.join(experiment_path, "fixedPointVSOdeX1.png"),
                       figname2=os.path.join(experiment_path, "fixedPointVSOdeX2.png"), useLogX=False)
        fitComparePlot(X1, X2, output, fitOutput, courbs,
                       figname=os.path.join(experiment_path, "fixedPointVSOdelogX1.png"),
                       figname2=os.path.join(experiment_path, "fixedPointVSOdelogX2.png"), useLogX=True)
        fitComparePlot(X1, X2, output, oldFitOutput, courbs,
                       figname=os.path.join(experiment_path, "oldfixedPointVSOdeX1.png"),
                       figname2=os.path.join(experiment_path, "oldfixedPointVSOdeX2.png"), useLogX=False)
        fitComparePlot(X1, X2, output, oldFitOutput, courbs,
                       figname=os.path.join(experiment_path, "oldfixedPointVSOdelogX1.png"),
                       figname2=os.path.join(experiment_path, "oldfixedPointVSOdelogX2.png"), useLogX=True)

        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(figsize=(19.2,10.8), dpi=100)
        cmap = plt.get_cmap('Dark2',X1.shape[0]*len(courbs))
        for idx1,x1 in enumerate(X1):
            for idx2,x2 in enumerate(X2):
                if idx2 in courbs and idx1==idx2:
                    ax.plot(testOfCp*kd,np.abs(styleFit[idx1,idx2,:])*kd,c=cmap(idx1*(courbs.index(idx2)+1)))
                    ax.plot(testOfCp*kd,np.abs(oldStyleFit[idx1,idx2,:]),c=cmap(idx1*(courbs.index(idx2)+1)),linestyle="--")
                    ax.axvline(competitions[idx1,idx2]*kd,c=cmap(idx1*(courbs.index(idx2)+1)),marker="o")
                    #ax.axvline(oldCompetitions[idx1,idx2],c=cmap(idx1*(courbs.index(idx2)+1)),marker="o",linestyle="--")
                    ax.axvline(kd,c=cmap(idx1*(courbs.index(idx2)+1)),marker="x",linestyle="-")
                    # ax.scatter(cpFitted[idx1,idx2],func(cpFitted[idx1,idx2],kd,Cactiv,CInhib,Kactiv,Kinhib,masks,myX_test[idx1,idx2]),c=cmap(idx1*(courbs.index(idx2)+1)),marker="x")
        ax.set_xscale("log")
        ax.set_yscale("log")
        fig.savefig("templateModelTwoLayer/equiNetwork_10_newfit/formOfcp.png")
        plt.show()
        print(competitions)

        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(figsize=(19.2,10.8), dpi=100)
        cmap = plt.get_cmap('Dark2',X1.shape[0]*len(courbs))
        for idx1,x1 in enumerate(X1):
            for idx2,x2 in enumerate(X2):
                if idx2 in courbs and idx1==idx2:
                    # ax.plot(testOfCp,np.abs(styleFit[idx1,idx2,:]),c=cmap(idx1*(courbs.index(idx2)+1)))
                    ax.plot(testOfCpt*kd,np.abs(visuCPT[idx1,idx2,:])+10**(-10),c=cmap(idx1*(courbs.index(idx2)+1)),linestyle="--")
                    # ax.axvline(competitions[idx1,idx2],c=cmap(idx1*(courbs.index(idx2)+1)),marker="o")
                    ax.axvline(outCPT[idx1,idx2]*kd,c=cmap(idx1*(courbs.index(idx2)+1)),marker="o",linestyle="--")
                    # ax.scatter(cpFitted[idx1,idx2],func(cpFitted[idx1,idx2],kd,Cactiv,CInhib,Kactiv,Kinhib,masks,myX_test[idx1,idx2]),c=cmap(idx1*(courbs.index(idx2)+1)),marker="x")
        ax.set_xscale("log")
        ax.set_yscale("log")
        fig.savefig("templateModelTwoLayer/equiNetwork_10_newfit/formOfcp2.png")
        plt.show()
        print(outCPT)


# 2.4: Write code to partition a linked list around a value x, such that all
# nodes less than x come before all nodes greater than or equeal to x. The
# order of the values does not matter so long as values less than x come before
# those that are greater than or equal to x.

# Solution: p. 212

from typing import Optional

class ListNode:
    def __init__(self, val: int = 0, next: Optional['ListNode'] = None):
        self.val = val
        self.next = next

def partition(head: 'ListNode', x: int) -> 'ListNode':
    point = head
    head = last = None
    right = None

    while point:
        if point.val >= x:
            temp = point
            point = point.next
            temp.next = right
            right = temp
        else:
            if not last:
                last = point

            temp = point
            point = point.next
            temp.next = head
            head = temp

    if last:
        last.next = right
        return head
    else:
        return right


"""Utils."""

import logging

from py_console import console, textColor

LEVEL_TO_COLOR = {
    logging.CRITICAL: textColor.MAGENTA,
    logging.ERROR: textColor.RED,
    logging.WARNING: textColor.YELLOW,
    logging.INFO: textColor.BLUE,
    logging.DEBUG: textColor.GREEN,
    logging.NOTSET: textColor.WHITE,
}

class CustomLoggingFormatter(logging.Formatter):
    def format(self, record):
        color = LEVEL_TO_COLOR[record.levelno]
        return console.highlight(
            f'({record.name}-{record.levelname}): {record.msg}', color
        )

class Log(logging.Logger):
    def __init__(self, name: str, level: int = logging.DEBUG):
        super(Log, self).__init__(name, level)
        self.propagate = False

        formatter = CustomLoggingFormatter()
        sh = logging.StreamHandler()
        sh.setLevel(logging.DEBUG)
        sh.setFormatter(formatter)
        self.handlers = [sh]


"""
Given two strings s and t, determine if they are isomorphic.

Two strings s and t are isomorphic if the characters in s can be replaced to get t.

All occurrences of a character must be replaced with another character while preserving the order of characters.
No two characters may map to the same character, but a character may map to itself.
"""

class Solution:
    def isIsomorphic(self, s: str, t: str) -> bool:
        s2t, t2s = {}, {}
        for letter_s, letter_t in zip(s, t):
            if letter_s in s2t and s2t[letter_s] != letter_t:
                return False

            if letter_t in t2s and t2s[letter_t] != letter_s:
                return False

            s2t[letter_s] = letter_t
            t2s[letter_t] = letter_s

        return True

if __name__ == '__main__':
    s = Solution()
    print(s.isIsomorphic("egg", "ads"))
    print(s.isIsomorphic("badc", "baba"))
    print(s.isIsomorphic("paper", "title"))


"""remove rotted column

Revision ID: 7e3df200a38d
Revises: c41abb284c8f
Create Date: 2021-10-20 17:16:08.299404

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = "7e3df200a38d"
down_revision = "c41abb284c8f"
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("weblinks", "rotted")
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "weblinks",
        sa.Column("rotted", sa.VARCHAR(), nullable=False, server_default="no"),
    )
    # ### end Alembic commands ###


# https://leetcode.com/problems/01-matrix/description/
# MEDIUM
# Tags: matrixlc, bfslc, dplc, #542

# GIVEN:
    # an m x n binary matrix, mat

# TASK:
    # return the distance of the nearest 0 for each cell
    # The distance between two adjacent cells is 1

# EXAMPLES:
    # Input: mat = [[0,0,0],
                #   [0,1,0],
                #   [0,0,0]]
    # Output: [[0,0,0],
            #  [0,1,0],
            #  [0,0,0]]

    # Input: mat = [[0,0,0],
                #   [0,1,0],
                #   [1,1,1]]
    # Output: [[0,0,0],
            #  [0,1,0],
            #  [1,2,1]]

###########################################################################################################

# ✅ ALGORITHM 1: BFS (WITH VISITED SET i.e. not space optimized)
# Create result matrix of the same dimensions as input matrix
# Create queue and hashset, visited
# For every element in matrix,
    # if element = 0, add it to queue, and add it to visited hashset
        # any q[i] = (r, c, distance_to_nearest_0)
        # we add 0's to visited as we don't need to visit them again (when they become neighbors of some cell), since their distance to nearest 0 is 0
# While queue is not empty, pop current element (and distance from 0) from queue
    # if current element's neighbor is not visited, increment current neighbor's distance from 0 by 1, and in result matrix, update current neighbor's cell with its updated distance-from-0 value
    # add current neighbor to visited
    # add current neighbor to queue, with incremented distance from 0
# Return result matrix

# TIME COMPLEXITY: O(m * n), where m = height of matrix, n = length of matrix
# SPACE COMPLEXITY: O(m * n)
    # result array is m * n size
    # set and queue are each approx. m*n size

def updateMatrix(mat):
    rows, cols = len(mat), len(mat[0])
    result = [[0] * len(row) for row in mat]
    q = []
    visited = set()

    for r in range(len(mat)):
        for c in range(len(mat[r])):
            if mat[r][c] == 0:
                q.append((r, c, 0))
                visited.add((r, c))

    # Perform bfs to update distance from 0 for each neighbor node
    while q:
        r, c, dist = q.pop(0)

        for x, y in [(0, -1), (0, 1), (-1, 0), (1, 0)]: # 4 neighbors, 1 in each direction
            nr, nc = r+x, c+y # index row and col of current neighbor
            if 0 <= nr < rows and 0 <= nc < cols and (nr, nc) not in visited: # if current neighbor is within matrix and it is not visited,
                result[nr][nc] = dist + 1 # update current neighbor's distance-from-0 value; since mat(nr,nc) is an immediate neighbor of mat(r,c), the distance-from-0 value of mat(nr, nc) = distance-from-0 value of mat(r,c) + 1
                visited.add((nr, nc))
                q.append((nr, nc, dist + 1)) # add current neighbor to queue, with its updated distance-from-0 value

    return result

#==========================================================================================================

# ✅ ALGORITHM 2: BFS (WITHOUT VISITED SET i.e. space optimized)
# Same as above, except we are modifying the matrix in-matrix (i.e. no additional result matrix to return)
# we mark out the 1's until they are updated; everything else (that are not marked) are already updated by previous iterations
# We also don't need a 3rd element in each tuple in the queue (i.e. the distance-from-0 value) to keep track of it

# TIME COMPLEXITY: O(m * n), where m = height of matrix, n = length of matrix
# SPACE COMPLEXITY: O(m * n)
    # length of queue might be m*n in the worst case scenario (if matrix contains all 0's)
    # but space complexity is still much better than above as we don't create visited hashset and result matrix

def updateMatrix(mat):
    rows, cols = len(mat), len(mat[0])
    q = []

    for r in range(len(mat)):
        for c in range(len(mat[r])):
            if mat[r][c] == 0:
                q.append((r, c))
            else:
                mat[r][c] = -1 # mark out all cells that require updating; all other cells (except the -1) are cells that are already updated with distance-from-0 values

    # Perform bfs to update distance from 0 for each neighbor node
    while q:
        r, c = q.pop(0)

        for x, y in [(0, -1), (0, 1), (-1, 0), (1, 0)]: # 4 neighbors, 1 in each direction
            nr, nc = r+x, c+y # index row and col of current neighbor
            if 0 <= nr < rows and 0 <= nc < cols and mat[nr][nc] == -1: # if current neighbor is within matrix and it is waiting to get updated
                mat[nr][nc] = mat[r][c] + 1 # update current neighbor's distance-from-0 value
                q.append((nr, nc)) # add current neighbor to queue, with its updated distance-from-0 value

    return mat

#==========================================================================================================

# ✅ ALGORITHM 3: DYNAMIC PROGRAMMING
# MAIN IDEA:
    # 1) traverse the whole matrix from the element in the top-left corner to the element in the bottom-right corner, checking the element above and left of each "1" cell
        # The minimum distance-to-0 value of current cell = 1 + the minimum of the distance-to-0 values of the cells above and to the left of current cell
            # if current cell is at the border and doesn't have a cell above/left of it, we set above/left as infinity so that it will not become the value assigned to current cell when comparing above and left cells for the minimum of the 2
    # 2) traverse the whole matrix in the reverse direction, i.e., from the element in the bottom-right corner to the element in the top-left corner, checking the element below and right of each non-zero cell
        # The minimum distance-to-0 value of the current cell = 1 + the minimum of the distance-to-0 values of the cells below and to the right of current cell
            # if current cell is at the border and doesn't have a cell below/right of it, we set below/right as infinity so that it will not become the value assigned to current cell when comparing below and right cells for the minimum of the 2
        # Then, we compare this minimum result with the previously stored minimum result in current cell, updating it with the smaller of the 2 values
    # 3) We return the matrix with the updated distance-to-0 values

# TIME COMPLEXITY: O(m * n), where m = height of matrix, n = length of matrix
# SPACE COMPLEXITY: O(1)
    # no extra space is needed

def updateMatrix(mat):
    rows, cols = len(mat), len(mat[0])

    # traverse the whole matrix from the element in the top-left corner to the element in the bottom-right corner, checking the element above and left of each "1" cell
    for r in range(rows):
        for c in range(cols):
            if mat[r][c] == 1:
                above = mat[r-1][c] if r > 0 else float('inf')
                left = mat[r][c-1] if c > 0 else float('inf')
                mat[r][c] = 1 + min(above, left)

    # traverse the whole matrix in the reverse direction, i.e., from the element in the bottom-right corner to the element in the top-left corner, checking the element below and right of each non-zero cell
    for r in range(rows-1, -1, -1):
        for c in range(cols-1, -1, -1):
            if mat[r][c] != 0:
                below = mat[r+1][c] if r < rows-1 else float('inf')
                right = mat[r][c+1] if c < cols-1 else float('inf')
                mat[r][c] = min(mat[r][c], below + 1, right + 1) # we compare the minimum between right and below, with previously stored minimum result in current cell, updating it with the smaller of the 2 values

    return mat


# Aaron Barnett

import numpy as np

class Perceptron:
    def __init__(self):
        np.random.seed(10)
        self.weights_h = np.random.uniform(low=-1, high=1, size=(784, 2))
        self.weights_o = np.random.uniform(low=-1, high=1, size=(2, 1))
        self.bias_h = np.random.uniform(low=-1, high=1, size=(2,))
        self.bias_o = np.random.uniform(low=-1, high=1, size=(1,))
        self.alpha = 0.5
        self.epochs = 1

    def train(self, data):
        for i in range(self.epochs):
            for j in range(len(data)):
                # perform backprop for output
                x = data[j]
                output = self.forward_prop(x)
                raw_error = data[j, 0] - output
                delta_output = raw_error * output * (1 - output)
                # backprop error to hidden layer
                input_layer = x[1:]
                hidden_layer = self.activate_layer(np.dot(self.weights_h.transpose(), input_layer) + self.bias_h)
                delta_hidden = np.multiply(np.multiply(hidden_layer, np.subtract(1, hidden_layer)), np.dot(self.weights_o, delta_output))
                # update output and internal weights
                for i in range(len(self.weights_o)):
                        self.weights_o[i] = self.weights_o[i] + self.alpha * hidden_layer[i] * delta_output
                for row in range(len(self.weights_h)):
                    for col in range(len(self.weights_h[row])):
                        self.weights_h[row, col] = self.weights_h[row, col] + self.alpha * input_layer[row] * delta_hidden[col]
                # update bias weights
                self.bias_o[0] = self.bias_o[0] + self.alpha * delta_output
                for row in range(len(self.bias_h)):
                    self.bias_h[row] = self.bias_h[row] + self.alpha * delta_hidden[row]

    def forward_prop(self, x):
        # x is 1x785 flat image with class label at index 0 (original 28x28 image)
        input_layer = x[1:]
        hidden_layer = self.activate_layer(np.dot(self.weights_h.transpose(), input_layer) + self.bias_h)
        output = self.activate_layer(np.dot(self.weights_o.transpose(), hidden_layer) + self.bias_o)
        return output

    def activate_layer(self, layer):
        # activation uses sigmoid
        activation = np.zeros(np.shape(layer))
        for i in range(len(layer)):
            activation[i] = 1/(1 + np.exp(-layer[i]))
        return activation

    def accuracy(self, data):
        correct = 0
        for row in range(len(data)):
            prediction = self.forward_prop(data[row])
            prediction = round(prediction[0])
            if prediction == data[row, 0]:
                correct += 1
        return correct/len(data)

def test_perceptron():
    train_data = np.genfromtxt('data/mnist_train_0_1.csv', delimiter=',')
    y = train_data[:, 0:1]
    train_data = np.divide(train_data[:, 1:], 255)
    train_data = np.c_[y, train_data]
    nn = Perceptron()
    nn.train(train_data)

    test_data = np.genfromtxt('data/mnist_test_0_1.csv', delimiter=',')
    y = test_data[:, 0:1]
    test_data = np.divide(test_data[:, 1:], 255)
    test_data = np.c_[y, test_data]
    print(nn.accuracy(test_data))

test_perceptron()


import os
import json
import time
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data.sampler import WeightedRandomSampler
import torchvision.transforms as transforms

from net.AutoEncoder import AutoEncoder, ConvAutoEncoder
from dataloader.cifar import CustomCifar
from dataloader.dataset import CustomDataset
from utils import conv_ae_train

transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomCrop(32, padding=5),
    transforms.ToTensor(),
    # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

cifar10 = CustomCifar()
cifar10.get_unbalanced_dataset()

trainset = CustomDataset(
    cifar10.train_data, cifar10.train_targets, transform=transform)

training_prefix = '20211111_2_conv_da_bn'

EPOCHS = 300
LR = 0.001
BATCH_SIZE = 128
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

save_dir = f'./ckpt/ae/{training_prefix}_Adam_{LR}_MSE_batch{BATCH_SIZE}_{EPOCHS}epoch_div8_unbalanced/'

# To deal with the unbalnced classes here we use the weighted random sampler.
weights = []
for idx in trainset.targets:
  if idx in [2, 4, 9]:
    weights.append(0.2)
  else:
    weights.append(0.1)

weights = torch.DoubleTensor(weights)
generator = torch.Generator(device='cpu')
generator.manual_seed(1)
sampler = WeightedRandomSampler(
    weights,
    len(trainset.targets),
    replacement=True,
    generator=generator)

trainloader = torch.utils.data.DataLoader(
    trainset,
    batch_size=BATCH_SIZE,
    shuffle=False,
    sampler=sampler,
    num_workers=2,
    drop_last=False)

# net = AutoEncoder()
net = ConvAutoEncoder()
loss_function = nn.MSELoss()
optimizer = optim.Adam(
    net.parameters(),
    lr=LR,
    weight_decay=1e-5)

os.makedirs(save_dir, exist_ok=True)

losses = []

print("Training Start")

for epoch in range(EPOCHS):
  start = time.time()

  avg_loss = conv_ae_train(
      trainloader,
      ae_model=net,
      loss_func=loss_function,
      optim=optimizer,
      # type='stacking',
      type='conv',
      mode='train')

  duration = time.time() - start

  for name, param in net.named_parameters():
    print(name, param.grad.abs().sum())

  losses.append(avg_loss)

  print(f"Epoch {epoch + 1}  Loss:{avg_loss} {round(duration,2)} sec")

  if (epoch + 1) % 10 == 0:
    torch.save(
        net.state_dict(),
        save_dir + f'model_{(epoch + 1) // 10}.pth'
    )

    loss_dict = {'loss': losses}
    with open(save_dir + 'metrics.json', 'w') as f:
      json.dump(loss_dict, f)


import sys
import heapq

n = int(sys.stdin.readline().rstrip())  # 입력에서 카드의 개수를 읽어옵니다.
cards = []  # 카드를 저장할 리스트
heapq.heapify(cards)  # cards 리스트를 최소 힙으로 만듭니다.
result = 0  # 결과 변수 (카드 섞는데 필요한 비용의 총합을 저장)

for i in range(n):
    card = int(sys.stdin.readline().rstrip())  # 카드의 가격을 읽어옵니다.
    heapq.heappush(cards, card)  # 카드의 가격을 최소 힙에 추가합니다.

flag = False  # 카드를 섞을 때 사용할 플래그 변수
while len(cards) != 1:
    mixing = 0  # 현재 섞인 카드의 가격을 저장할 변수
    if not flag:
        for j in range(2):
            a = heapq.heappop(cards)  # 가장 가격이 낮은 카드 2장을 꺼냅니다.
            mixing += a  # 두 장의 카드를 섞은 가격을 누적합니다.
        flag = True  # 플래그를 True로 설정하여 다음 루프에서 1장의 카드를 꺼내야 합니다.
    else:
        a = heapq.heappop(cards)  # 가장 가격이 낮은 카드 1장을 꺼냅니다.
        b = heapq.heappop(cards)  # 가장 가격이 낮은 카드 1장을 또 꺼냅니다.
        mixing += (a + b)  # 두 장의 카드를 섞은 가격을 누적합니다.

    result += mixing  # 섞은 카드의 가격을 결과에 더합니다.
    heapq.heappush(cards, mixing)  # 섞은 카드의 가격을 최소 힙에 추가합니다.

# 결과 출력: 모든 카드를 섞고 남은 한 장의 카드의 가격
print(result)


import pytest
from flask import template_rendered

@pytest.mark.parametrize(
    ("path", "template_name"),
    (
        ("/", "xhr.html"),
        ("/plain", "xhr.html"),
        ("/fetch", "fetch.html"),
        ("/jquery", "jquery.html"),
    ),
)
def test_index(app, client, path, template_name):
    def check(sender, template, context):
        assert template.name == template_name

    with template_rendered.connected_to(check, app):
        client.get(path)

@pytest.mark.parametrize(
    ("a", "b", "result"), ((2, 3, 5), (2.5, 3, 5.5), (2, None, 2), (2, "b", 2))
)
def test_add(client, a, b, result):
    response = client.post("/add", data={"a": a, "b": b})
    assert response.get_json()["result"] == result


""""
Task 3 - Football Cards
"""
cards = input().split(" ")
teama = ["A-1", "A-2", "A-3", "A-4", "A-5", "A-6", "A-7", "A-8", "A-9", "A-10", "A-11"]
teamb = ["B-1", "B-2", "B-3", "B-4", "B-5", "B-6", "B-7", "B-8", "B-9", "B-10", "B-11"]

if cards!="":
    for card in cards:
        if card in teama:
            teama.remove(card)
        elif card in teamb:
            teamb.remove(card)
        if len(teama) < 7 or len(teamb) < 7:
            break

print(f"Team A - {len(teama)}; Team B - {len(teamb)}")
if len(teama) < 7 or len(teamb) < 7:
    print("Game was terminated")


import json

import numpy as np
import torch
from torch.utils.data import Dataset
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, confusion_matrix
import logging
import os
from .tokenizers import *

logger = logging.getLogger(os.path.basename(__file__))

class BaseDataset(Dataset):
    def __init__(self, args, n_samples, n_classes, transform, doc_index):
        self.args = args
        self.n_samples = n_samples
        self.n_classes = n_classes
        self.transform = transform
        self.doc_index = doc_index

    def __len__(self):
        return self.n_samples

    @staticmethod
    def metric(inputs, labels, logits, test_mode=True, case_study=False):
        """
        logits and labels are numpy if test_mode is True else torch
        Args:
            inputs:
            labels:
            logits:
            test_mode:

        Returns:

        """
        labels = labels[0]
        logits = logits[0]
        metric = {}
        if test_mode:
            pred_labels = np.argmax(logits, axis=-1)
            acc = accuracy_score(labels, pred_labels)
            precision = precision_score(labels, pred_labels, average='macro')
            recall = recall_score(labels, pred_labels, average='macro')
            f1 = f1_score(labels, pred_labels, average='macro')
            cf = confusion_matrix(labels, pred_labels)
            cf = json.dumps(cf.tolist())
            logger.info(f'confusion matrix\n{cf}')
            metric.update({'acc': acc, 'precision': precision, 'recall': recall, 'f1': f1})
            for key, val in metric.items():
                logger.info(f'{key}: {val:.4f}')
        else:
            score, pred_labels = logits.max(-1)
            acc = float((pred_labels == labels).long().sum()) / labels.size(0)
            metric.update({'acc': acc})
        return metric

class NLPDataset(BaseDataset):
    def __init__(self, args, n_samples, n_classes=None, transform=None, doc_index=None):
        super(NLPDataset, self).__init__(args, n_samples, n_classes, transform, doc_index)
        self.text = args['text']
        self.doc = args['doc']

        if isinstance(args['input_ids'][0][0], int):
            self.input_ids = torch.LongTensor(args['input_ids'])
        else:
            self.input_ids = torch.FloatTensor(args['input_ids'])
        if 'attention_mask' in args:
            self.attention_mask = torch.LongTensor(args['attention_mask'])
        else:
            self.attention_mask = None
        self.label = torch.LongTensor(args['label'])

    def __getitem__(self, item):
        if self.attention_mask is not None:
            data = [self.input_ids[item], self.attention_mask[item]]
        else:
            data = [self.input_ids[item]]
        label = [self.label[item]]
        return data, label

class NERDataset(NLPDataset):
    def __init__(self, args, n_samples, n_classes=None, transform=None, doc_index=None):
        super(NERDataset, self).__init__(args, n_samples, n_classes, transform, doc_index)
        self.label_vocab = args['label vocab']

    def metric(self, inputs, labels, logits, test_mode=True, case_study=False):
        """
        logits and labels are numpy if test_mode is True else torch
        Args:
            inputs:
            labels:
            logits:
            test_mode:

        Returns:

        """
        labels = labels[0]  # (B, L)
        logits = logits[0]  # (B, L, C)
        metric = {}
        if test_mode:
            pred_labels = np.argmax(logits, axis=-1)  # (B, L)
            labels, pred_labels = labels.flatten(), pred_labels.flatten()
            acc = accuracy_score(labels, pred_labels)
            precision = precision_score(labels, pred_labels, average='macro')
            recall = recall_score(labels, pred_labels, average='macro')
            f1 = f1_score(labels, pred_labels, average='macro')
            cf = confusion_matrix(labels, pred_labels)
            cf = json.dumps(cf.tolist())
            logger.info(f'confusion matrix\n{cf}')
            metric.update({'acc': acc, 'precision': precision, 'recall': recall, 'f1': f1})
            for key, val in metric.items():
                logger.info(f'{key}: {val:.4f}')
            true_mask = labels != self.label_vocab['O']
            true_acc = float(((pred_labels == labels) & true_mask).sum() / true_mask.sum())
            logger.info(f'true acc: {true_acc:.4f}')
        else:
            _, pred_labels = logits.max(-1)
            labels, pred_labels = labels.flatten(), pred_labels.flatten()
            acc = float((pred_labels == labels).long().sum()) / labels.size(0)
            metric.update({'acc': acc})
            true_mask = labels != self.label_vocab['O']
            true_acc = float(((pred_labels == labels) & true_mask).sum() / true_mask.sum())
            metric.update({'true acc': true_acc})
        return metric

class ImageDataset(BaseDataset):
    def __init__(self, args, n_classes=None, transform=None, doc_index=None):
        super(ImageDataset, self).__init__(args, n_classes, transform, doc_index)
        self.data = args[0]
        self.targets = torch.LongTensor(args[1])

    def __getitem__(self, item):
        data = [self.transform(self.data[item])]
        target = [self.targets[item]]
        return data, target

class ImageCCSADataset(BaseDataset):
    def __init__(self, args, n_classes=None, transform=None, doc_index=None):
        super(ImageCCSADataset, self).__init__(args, n_classes, transform, doc_index)
        num_sample = 5
        self.src_x = args[0]
        self.tgt_x = args[1]
        self.src_y = torch.LongTensor(args[2])
        self.tgt_y = torch.LongTensor(args[3])

        tgt_idx = [np.where(self.tgt_y == i)[0] for i in range(n_classes)]
        tgt_idx = [np.random.choice(idx, num_sample) for idx in tgt_idx]
        tgt_idx = np.concatenate(tgt_idx)
        self.tgt_x = self.tgt_x[tgt_idx]
        self.tgt_y = self.tgt_y[tgt_idx]

        positive, negative = [], []

        for trs in range(len(self.src_y)):
            for trt in range(len(self.tgt_y)):
                if self.src_y[trs] == self.tgt_y[trt]:
                    positive.append([trs, trt])
                else:
                    negative.append([trs, trt])
        logger.info(f"num of positive/negative pairs: {len(positive)}/{len(negative)}")
        np.random.shuffle(negative)
        self.pairs = positive + negative[:3 * len(positive)]
        random.shuffle(self.pairs)
        self.n_samples = len(self.pairs)

    def __getitem__(self, item):
        src_idx, tgt_idx = self.pairs[item]
        data = [self.transform(self.src_x[src_idx]), self.transform(self.tgt_x[tgt_idx])]
        targets = [self.src_y[src_idx], self.tgt_y[tgt_idx]]
        return data, targets

    @staticmethod
    def metric(inputs, labels, logits, test_mode=True):
        """
        logits and labels are numpy if test_mode is True else torch
        Args:
            inputs:
            labels:
            logits:
            test_mode:

        Returns:

        """
        src_labels, tgt_labels = labels
        src_logits, tgt_logits = logits
        metric = {}
        src_score, src_pred_labels = src_logits.max(-1)
        src_acc = float((src_pred_labels == src_labels).long().sum()) / src_labels.size(0)
        tgt_score, tgt_pred_labels = tgt_logits.max(-1)
        tgt_acc = float((tgt_pred_labels == tgt_labels).long().sum()) / tgt_labels.size(0)

        metric.update({'src_acc': src_acc, 'tgt_acc': tgt_acc})
        return metric

class REDataset(BaseDataset):
    def __init__(self, args, n_samples, n_classes=None, transform=None, doc_index=None):
        super(REDataset, self).__init__(args, n_samples, n_classes, transform, doc_index)
        self.text = args['text']
        self.tokens = args['tokens']
        self.doc = args['doc']

        self.input_ids = torch.LongTensor(args['input_ids'])
        self.attention_mask = torch.LongTensor(args['attention_mask'])
        self.e1_mask = torch.LongTensor(args['e1_mask'])
        self.e2_mask = torch.LongTensor(args['e2_mask'])
        self.label = torch.LongTensor(args['label'])

        zero = torch.zeros_like(self.input_ids)
        self.mlm_input_ids = torch.LongTensor(args['mlm_input_ids']) if 'mlm_input_ids' in args else zero
        self.mlm_labels = torch.LongTensor(args['mlm_labels']) if 'mlm_labels' in args else zero
        self.valid_ids = torch.LongTensor(args['valid_ids']) if 'valid_ids' in args else zero
        # self.dep_matrix = torch.LongTensor(args['dep_matrix']) if 'dep_matrix' in args else zero
        if 'dep_text' in args:
            self.dep_text = args['dep_text']
        self.dep_matrix = zero

    def __getitem__(self, item):
        data = [self.input_ids[item], self.attention_mask[item], self.e1_mask[item], self.e2_mask[item],
                self.valid_ids[item], self.mlm_input_ids[item], self.dep_matrix[item]]
        label = [self.label[item], self.mlm_labels[item]]
        return data, label

    def metric(self, inputs, labels, logits, test_mode=True, case_study=False):
        """
        logits and labels are numpy if test_mode is True else torch
        Args:
            case_study:
            inputs:
            labels:
            logits:
            test_mode:

        Returns:

        """
        # input_ids, attention_mask, e1_mask, e2_mask, valid_ids, mlm_input_ids, dep_matrix = inputs
        re_labels, mlm_labels = labels
        re_logits = logits[0]

        metric = {}
        if test_mode:
            re_pred_labels = np.argmax(re_logits, axis=-1)
            acc = accuracy_score(re_labels, re_pred_labels)
            precision = precision_score(re_labels, re_pred_labels, average='macro')
            recall = recall_score(re_labels, re_pred_labels, average='macro')
            f1 = f1_score(re_labels, re_pred_labels, average='macro')
            cf = confusion_matrix(re_labels, re_pred_labels)
            cf = json.dumps(cf.tolist())
            logger.info(f'confusion matrix\n{cf}')
            metric.update({'acc': acc, 'precision': precision, 'recall': recall, 'f1': f1})
            for key, val in metric.items():
                logger.info(f'{key}: {val:.4f}')
            # if np.max(mlm_labels) > 0:
            #     _, mlm_logits, unmask_logits = logits
            #     n1, n2 = len(mlm_logits), len(mlm_logits[0])
            #     score = np.zeros((n1, n2))
            #     for i in range(n1):
            #         for j, (logits, label, input_id) in enumerate(zip(mlm_logits[i], mlm_labels[i], mlm_input_ids[i])):
            #             logits = np.exp(logits)
            #             logits /= np.sum(logits, axis=-1)
            #             if label != -100:
            #                 score[i][j] = logits[label]
            #             else:
            #                 score[i][j] = logits[input_id]
            #
            #     unmask_labels = input_ids.copy()
            #     unmask_labels[(mlm_input_ids == 103) | (input_ids == 0)] = -100
            #
            #     mlm_pred_labels = np.argmax(mlm_logits, axis=-1)
            #     mlm_mask = (mlm_labels != -100)
            #     mlm_acc = ((mlm_pred_labels == mlm_labels) & mlm_mask).sum() / mlm_mask.sum()
            #     metric.update({'mlm acc': mlm_acc})
            #     logger.info(f'mlm acc: {mlm_acc:.4f}')
            #
            #     unmask_pred_labels = np.argmax(unmask_logits, axis=-1)
            #     unmask_mask = (unmask_labels != -100)
            #     unmask_acc = ((unmask_pred_labels == unmask_labels) & unmask_mask).sum() / unmask_mask.sum()
            #     metric.update({'unmask acc': unmask_acc})
            #     logger.info(f'unmask acc: {unmask_acc:.4f}')
            if case_study:
                self.case_study(re_pred_labels, re_labels, re_logits, score, mlm_labels)
        else:
            score, re_pred_labels = re_logits.max(-1)
            acc = float((re_pred_labels == re_labels).long().sum()) / re_labels.size(0)
            metric.update({'acc': acc})

            # if torch.max(mlm_labels) > 0:
            #     _, mlm_logits, unmask_logits = logits
            #     unmask_labels = input_ids.masked_fill((mlm_input_ids == 103) | (input_ids == 0), -100)
            #
            #     score, mlm_pred_labels = mlm_logits.max(-1)
            #     mlm_mask = (mlm_labels != -100)
            #     mlm_acc = float(((mlm_pred_labels == mlm_labels) & mlm_mask).sum() / mlm_mask.sum())
            #     metric.update({'mlm acc': mlm_acc})
            #
            #     score, unmask_pred_labels = unmask_logits.max(-1)
            #     unmask_mask = (unmask_labels != -100)
            #     unmask_acc = float(((unmask_pred_labels == unmask_labels) & unmask_mask).sum() / unmask_mask.sum())
            #     metric.update({'unmask acc': unmask_acc})

        return metric

    def case_study(self, pre_labels, labels, logits, score, mlm_labels):
        logits = np.exp(logits)
        logits /= np.sum(logits, axis=1, keepdims=True)
        for tc in range(self.n_classes):
            for fc in range(self.n_classes):
                if tc != fc:
                    loc = np.where((labels == tc) & (pre_labels == fc))[0]
                    for l in loc:
                        print(self.text[l] + f' {tc} {fc} {logits[l]}')
                        print(self.tokens[l])
                        print(score[l])
                        print(mlm_labels[l])

        # for tc in range(self.n_classes):
        #     loc = np.where(labels == tc)[0]
        #     bad_case += [self.text[l] + f'|{tc}' for l in loc]
        # for case in bad_case:
        #     print(case)

class REMDomainAdaptationDataset(REDataset):
    def __init__(self, args, n_samples, n_classes=None, transform=None, doc_index=None):
        super(REMDomainAdaptationDataset, self).__init__(args, n_samples, n_classes, transform, doc_index)
        self.doc = torch.LongTensor(args['doc'])

    def __getitem__(self, item):
        data = [self.input_ids[item], self.attention_mask[item], self.e1_mask[item], self.e2_mask[item],
                self.valid_ids[item], self.dep_matrix[item]]
        label = [self.label[item], self.doc[item]]
        return data, label

    def metric(self, inputs, labels, logits, test_mode=True, case_study=False):
        """
        logits and labels are numpy if test_mode is True else torch
        Args:
            inputs:
            labels:
            logits:
            test_mode:

        Returns:

        """
        re_labels, doc = labels
        re_logits, domain_logits = logits

        metric = {}
        if test_mode:
            re_pred_labels = np.argmax(re_logits, axis=-1)
            acc = accuracy_score(re_labels, re_pred_labels)
            precision = precision_score(re_labels, re_pred_labels, average='macro')
            recall = recall_score(re_labels, re_pred_labels, average='macro')
            f1 = f1_score(re_labels, re_pred_labels, average='macro')
            cf = confusion_matrix(re_labels, re_pred_labels)
            cf = json.dumps(cf.tolist())
            logger.info(f'confusion matrix\n{cf}')
            metric.update({'acc': acc, 'precision': precision, 'recall': recall, 'f1': f1})
            for key, val in metric.items():
                logger.info(f'{key}: {val:.4f}')

            domain_pred_labels = np.argmax(domain_logits, axis=-1)
            domain_acc = accuracy_score(doc, domain_pred_labels)
            metric.update({'domain acc': domain_acc})
            logger.info(f'domain acc: {domain_acc:.4f}')
            if case_study:
                self.case_study(re_pred_labels, re_labels)
        else:
            score, re_pred_labels = re_logits.max(-1)
            acc = float((re_pred_labels == re_labels).long().sum()) / re_labels.size(0)
            metric.update({'acc': acc})

            score, domain_pred_labels = domain_logits.max(-1)
            domain_acc = float((domain_pred_labels == doc).long().sum() / doc.size(0))
            metric.update({'domain acc': domain_acc})
        return metric

class RelationExtractionMMDTrainDataset(BaseDataset):
    def __init__(self, args, n_classes=None, transform=None, doc_index=None):
        super(RelationExtractionMMDTrainDataset, self).__init__(args, n_classes, transform, doc_index)
        self.data = args[0]
        self.pos1 = args[1]
        self.pos2 = args[2]
        self.docs = args[3]
        self.targets = torch.LongTensor(args[4])

        src_idx = np.where(self.docs == 0)[0]
        tgt_idx = np.where(self.docs == 1)[0]
        src_num = len(src_idx)
        tgt_num = len(tgt_idx)

        if src_num > tgt_num:
            ratio = src_num // tgt_num
            tmp = np.repeat(tgt_idx, ratio)
            tgt_idx = np.concatenate([tgt_idx, np.random.choice(tmp, src_num - tgt_num, replace=False)])
            np.random.shuffle(tgt_idx)
        else:
            ratio = tgt_num // src_num
            tmp = np.repeat(src_idx, ratio)
            src_idx = np.concatenate([src_idx, np.random.choice(tmp, tgt_num - src_num, replace=False)])
            np.random.shuffle(src_idx)

        self.n_samples = max(src_num, tgt_num)

        ls = [self.data, self.pos1, self.pos2, self.targets]
        self.src_data, self.src_pos1, self.src_pos2, self.src_targets = [x[src_idx] for x in ls]
        self.tgt_data, self.tgt_pos1, self.tgt_pos2, self.tgt_targets = [x[tgt_idx] for x in ls]

    def __getitem__(self, item):
        src = [self.src_data[item], self.src_pos1[item], self.src_pos2[item]]
        tgt = [self.tgt_data[item], self.tgt_pos1[item], self.tgt_pos2[item]]
        data = src + tgt
        targets = [self.src_targets[item], self.tgt_targets[item]]
        return data, targets

    @staticmethod
    def metric(inputs, labels, logits, test_mode=True):
        """
        logits and labels are numpy if test_mode is True else torch
        Args:
            inputs:
            labels:
            logits:
            test_mode:

        Returns:

        """
        src_labels, tgt_labels = labels
        src_logits, tgt_logits = logits
        metric = {}
        src_score, src_pred_labels = src_logits.max(-1)
        src_acc = float((src_pred_labels == src_labels).long().sum()) / src_labels.size(0)
        tgt_score, tgt_pred_labels = tgt_logits.max(-1)
        tgt_acc = float((tgt_pred_labels == tgt_labels).long().sum()) / tgt_labels.size(0)

        metric.update({'src_acc': src_acc, 'tgt_acc': tgt_acc})
        return metric

class RelationExtractionCCSADataset(BaseDataset):
    def __init__(self, args, n_classes=None, transform=None, doc_index=None):
        super(RelationExtractionCCSADataset, self).__init__(args, n_classes, transform, doc_index)
        self.data = args[0]
        self.pos1 = args[1]
        self.pos2 = args[2]
        self.docs = args[3]
        self.targets = torch.LongTensor(args[4])

        src_idx = np.where(self.docs == 0)[0]
        tgt_idx = np.where(self.docs == 1)[0]
        src_num = len(src_idx)
        tgt_num = len(tgt_idx)

        positive, negative = [], []

        self.src_y = self.targets[src_idx]
        self.tgt_y = self.targets[tgt_idx]
        for trs in range(len(self.src_y)):
            for trt in range(len(self.src_y)):
                if self.src_y[trs] == self.tgt_y[trt]:
                    positive.append([trs, trt])
                else:
                    negative.append([trs, trt])
        logger.info(f"num of positive/negative pairs: {len([positive])}/{len(negative)}")
        self.pairs = positive + negative
        random.shuffle(self.pairs)

    def __getitem__(self, item):
        src_idx, tgt_idx = self.pairs[item]
        data = [self.data[src_idx], self.pos1[src_idx], self.pos2[src_idx],
                self.data[tgt_idx], self.pos1[tgt_idx], self.pos2[tgt_idx]]
        targets = [self.targets[src_idx], self.targets[tgt_idx]]
        return data, targets

    @staticmethod
    def metric(inputs, labels, logits, test_mode=True):
        """
        logits and labels are numpy if test_mode is True else torch
        Args:
            inputs:
            labels:
            logits:
            test_mode:

        Returns:

        """
        src_labels, tgt_labels = labels
        src_logits, tgt_logits = logits
        metric = {}
        src_score, src_pred_labels = src_logits.max(-1)
        src_acc = float((src_pred_labels == src_labels).long().sum()) / src_labels.size(0)
        tgt_score, tgt_pred_labels = tgt_logits.max(-1)
        tgt_acc = float((tgt_pred_labels == tgt_labels).long().sum()) / tgt_labels.size(0)

        metric.update({'src_acc': src_acc, 'tgt_acc': tgt_acc})
        return metric


# strip out a set of nuisance html attributes that can mess up rendering in
# RSS feeds
import re
from lxml.html.clean import Cleaner

bad_attrs = ['width', 'height', 'style', '[-a-z]*color',
             'background[-a-z]*', 'on*', 'id']
single_quoted = "'[^']+'"
double_quoted = '"[^"]+"'
non_space = '[^ "\'>]+'
cstr = ("<"  # open
        "([^>]+) "  # prefix
        "(?:%s) *" % ('|'.join(bad_attrs),) +  # undesirable attributes
        '= *(?:%s|%s|%s)' % (non_space, single_quoted, double_quoted) +  # value
        "([^>]*)" +  # postfix
        ">")
htmlstrip = re.compile("<"  # open
                       "([^>]+) "  # prefix
                       "(?:%s) *" % ('|'.join(bad_attrs),) +  # undesirable attributes
                       '= *(?:%s|%s|%s)' % (non_space, single_quoted, double_quoted) +  # value
                       "([^>]*)"  # postfix
                       ">",       # end
                       re.I)

def clean_attributes(html_str):
    """移除HTML标签中无用的属性, 即上面的bad_attrs
    例如:
content

    变成:
content

    考虑连id和class一起移除
    """
    empty_tag = re.compile(r"<\w+[^>]*>[\s| ]*</\w+>", re.I)
    html_str = re.sub(empty_tag, '', html_str)
    while htmlstrip.search(html_str):
        html_str = htmlstrip.sub(r'<\1\2>', html_str)
    return html_str

def normalize_spaces(s):
    """replace any sequence of whitespace characters with a single space
    将多个空格变成一个空格
    """
    if not s:
        return ''
    return ' '.join(s.split())

def text_content(s):
    if not s:
        return ''
    return re.sub(r'</?\w+[^>]*>', '', s)

html_cleaner = Cleaner(scripts=True, javascript=True, comments=True,
                       style=True, links=True, meta=False, add_nofollow=True,
                       page_structure=False, processing_instructions=True,
                       embedded=False, frames=False, forms=False,
                       annoying_tags=False, remove_tags=None,
                       kill_tags=['footer', 'nav', 'input'],
                       remove_unknown_tags=False, safe_attrs_only=False)

# kill_tags 不能删除tag说明
# form: 部分网站直接将form作为最外层- -|||


import re

HTML_COLORS = [
    "aliceblue",
    "antiquewhite",
    "aqua",
    "aquamarine",
    "azure",
    "beige",
    "bisque",
    "black",
    "blanchedalmond",
    "blue",
    "blueviolet",
    "brown",
    "burlywood",
    "cadetblue",
    "chartreuse",
    "chocolate",
    "coral",
    "cornflowerblue",
    "cornsilk",
    "crimson",
    "cyan",
    "darkblue",
    "darkcyan",
    "darkgoldenrod",
    "darkgray",
    "darkgreen",
    "darkkhaki",
    "darkmagenta",
    "darkolivegreen",
    "darkorange",
    "darkorchid",
    "darkred",
    "darksalmon",
    "darkseagreen",
    "darkslateblue",
    "darkslategray",
    "darkturquoise",
    "darkviolet",
    "deeppink",
    "deepskyblue",
    "dimgray",
    "dodgerblue",
    "firebrick",
    "floralwhite",
    "forestgreen",
    "fuchsia",
    "gainsboro",
    "ghostwhite",
    "gold",
    "goldenrod",
    "gray",
    "green",
    "greenyellow",
    "honeydew",
    "hotpink",
    "indianred",
    "indigo",
    "ivory",
    "khaki",
    "lavender",
    "lavenderblush",
    "lawngreen",
    "lemonchiffon",
    "lightblue",
    "lightcoral",
    "lightcyan",
    "lightgoldenrodyellow",
    "lightgray",
    "lightgreen",
    "lightpink",
    "lightsalmon",
    "lightsalmon",
    "lightseagreen",
    "lightskyblue",
    "lightslategray",
    "lightsteelblue",
    "lightyellow",
    "lime",
    "limegreen",
    "linen",
    "magenta",
    "maroon",
    "mediumaquamarine",
    "mediumblue",
    "mediumorchid",
    "mediumpurple",
    "mediumseagreen",
    "mediumslateblue",
    "mediumslateblue",
    "mediumspringgreen",
    "mediumturquoise",
    "mediumvioletred",
    "midnightblue",
    "mintcream",
    "mistyrose",
    "moccasin",
    "navajowhite",
    "navy",
    "oldlace",
    "olive",
    "olivedrab",
    "orange",
    "orangered",
    "orchid",
    "palegoldenrod",
    "palegreen",
    "paleturquoise",
    "palevioletred",
    "papayawhip",
    "peachpuff",
    "peru",
    "pink",
    "plum",
    "powderblue",
    "purple",
    "rebeccapurple",
    "red",
    "rosybrown",
    "royalblue",
    "saddlebrown",
    "salmon",
    "sandybrown",
    "seagreen",
    "seashell",
    "sienna",
    "silver",
    "skyblue",
    "slateblue",
    "slategray",
    "snow",
    "springgreen",
    "steelblue",
    "tan",
    "teal",
    "thistle",
    "tomato",
    "turquoise",
    "violet",
    "wheat",
    "white",
    "whitesmoke",
    "yellow",
    "yellowgreen",
]

def validate_color(color: str) -> bool:
    match_hex = re.search(r"^#(?:[0-9a-f]{3}){1,2}$", color.lower())
    match_html_color_name = color.lower() in HTML_COLORS
    return bool(match_hex or match_html_color_name)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from setuptools import setup, find_packages
import glob
import ntpath

def get_module_name(module_path):
    """
    Return the module name of the module path
    """
    return ntpath.split(module_path)[1].split(".")[0]

def snake_to_camel(word):
    """
    Convert a word from snake_case to CamelCase
    """
    return ''.join(x.capitalize() or '_' for x in word.split('_'))

setup(
    name='fn_aws_iam',
    version='1.0.1',
    license='MIT',
    author='IBM Resilient',
    author_email='',
    url='https://ibm.com/mysupport',
    description="Amazon AWS IAM Integration for Resilient",
    long_description="Integration with Amazon AWS IAM to facilitate manual enrichment and targeted remediation actions. "
                     "Teams can investigate an attack by searching for AWS user accounts across an AWS environment, and "
                     "quickly respond to attacks by executing remediation actions, such as removing permissions or "
                     "login profiles for suspicious accounts from within the Resilient platform.",
    install_requires=[
        'resilient_circuits>=35.0.0',
        'resilient_lib>=35.0.0',
        'boto3>=1.10.6'
    ],
    packages=find_packages(),
    include_package_data=True,
    platforms='any',
    classifiers=[
        'Programming Language :: Python',
    ],
    entry_points={
        "resilient.circuits.components": [
            # When setup.py is executed, loop through the .py files in the components directory and create the entry points.
            "{}FunctionComponent = fn_aws_iam.components.{}:FunctionComponent".format(snake_to_camel(get_module_name(filename)), get_module_name(filename)) for filename in glob.glob("./fn_aws_iam/components/[a-zA-Z]*.py")
        ],
        "resilient.circuits.configsection": ["gen_config = fn_aws_iam.util.config:config_section_data"],
        "resilient.circuits.customize": ["customize = fn_aws_iam.util.customize:customization_data"],
        "resilient.circuits.selftest": ["selftest = fn_aws_iam.util.selftest:selftest_function"]
    }
)


from datetime import datetime
import enum
import gi
gi.require_version('Gtk', '3.0')

from gi.repository import Gtk, GObject

class FeedField(enum.IntEnum):
    TIMESTAMP = 0
    COMMAND = 1
    CONTENT = 2

    def __str__(self):
        return f'{self.name.replace("_", " ").title()}'

class FeedData(Gtk.ListStore):
    """Widget containing feed text"""
    def __init__(self):
        Gtk.ListStore.__init__(self, str, str, str)

        self._iters = {}
        self._content = {}
        self.last_msg = 0

    def append_message(self, date_time, command, content):
        timestamp = date_time.timestamp()
        time_str = FeedData.datetime_to_string(date_time)

        message = [time_str, f'{command}:', content]
        self._iters[timestamp] = self.append(message)
        self._content[timestamp] = (command, content)
        self.last_msg = len(self._content) - 1

    @classmethod
    def datetime_to_string(cls, date_time):
        second_str = float(f'{date_time:%S.%f}')
        return f'[{date_time:%H:%M:}{second_str:06.3f}]'

    def clear_all(self):
        self.clear()
        self._iters = {}
        self._content = {}
        self.last_msg = 0

class FeedTable(Gtk.ScrolledWindow):
    """The thing that shows a FeedData buffer for eyes"""
    def __init__(self, feed_data):
        Gtk.ScrolledWindow.__init__(self)

        self.view = Gtk.TreeView(feed_data)

        self.column_names = [str(field) for field in FeedField]
        self.columns = [None] * len(self.column_names)

        def init_column(col, field, cell_renderer=Gtk.CellRendererText):
            cell = cell_renderer()

            col.pack_start(cell, True)
            col.add_attribute(cell, 'text', field)

        for field in FeedField:
            self.columns[field] = Gtk.TreeViewColumn(self.column_names[field])
            init_column(self.columns[field], field)
            self.view.append_column(self.columns[field])

        self.view.set_headers_visible(False)

        selection = self.view.get_selection()
        selection.set_mode(Gtk.SelectionMode.NONE)

        self.set_min_content_width(420)
        self.set_min_content_height(100)

        self.set_propagate_natural_width(True)
        self.set_propagate_natural_height(True)

        self.add(self.view)

class FeedDisplay(Gtk.Frame):
    """Display widget for feed"""
    def __init__(self):
        Gtk.Frame.__init__(self, label='Feed')

        self.data = FeedData()
        self.table = FeedTable(self.data)

        self.grid = Gtk.Grid()

        self.grid.attach(self.table, 1, 1, 1, 1)

        self.grid.set_column_spacing(15)
        self.grid.set_row_spacing(5)

        self.grid.set_margin_start(5)
        self.grid.set_margin_end(5)
        self.grid.set_margin_top(5)
        self.grid.set_margin_bottom(5)

        self.set_margin_start(5)
        self.set_margin_end(5)
        self.set_margin_top(5)
        self.set_margin_bottom(5)

        self.add(self.grid)

        self.set_name('feed-display')

    def add_to_feed(self, command, content):
        self.data.append_message(datetime.now(), command, content)
        path = Gtk.TreePath.new_from_indices([self.data.last_msg])
        self.table.view.scroll_to_cell(path, None, True, 1.0, 0.0)


#peak_speed_all_long = np.reshape(peak_speed_all, (n_par, 2, n_trials*2))

# Compute significance for each sample
#t_all, p_all = stats.ttest_rel(peak_speed_all_long[:,0,:], peak_speed_all_long[:,1,:], axis=0)
#sig_samp = np.where(p_all < 0.05)[0]

def fill_outliers(array):
    """Fill outliers in 1D array using the mean of surrounding non-outliers"""
    # Get index of outliers
    idx_outlier = np.where(np.abs(zscore(array)) > 3)[0]
    idx_non_outlier = np.where(np.abs(zscore(array)) <= 3)[0]
    # Fill each outlier with mean of closest non outlier
    for idx in idx_outlier:
        # Get index of the closest non-outlier before and after
        where_before = np.where(idx_non_outlier < idx)[0]
        where_after = np.where(idx_non_outlier > idx)[0]
        if len(where_before) > 0 and len(where_after) > 0:  # Middle sample
            array[idx] = np.mean([array[idx_non_outlier[where_before[-1]]], array[idx_non_outlier[where_after[0]]]])
        elif len(where_before) == 0 and len(where_after) > 0:  # First sample
            array[idx] = np.mean([array[idx_non_outlier[where_after[1]]], array[idx_non_outlier[where_after[0]]]])
        elif len(where_before) > 0 and len(where_after) == 0:  # Last sample
            array[idx] = np.mean([array[idx_non_outlier[where_before[-2]]], array[idx_non_outlier[where_before[-1]]]])
    return array

def plot_conds(array, var=None, color_slow="#00863b", color_fast="#3b0086"):
    """array = (conds x blocks x trials)
    Plot data divided into two conditions, if given add the variance as shaded area"""
    # Plot without the first 5 movements
    plt.plot(array[0, :, :].flatten()[5:], label="Slow", color=color_slow, linewidth=3)
    plt.plot(array[1, :, :].flatten()[5:], label="Fast", color=color_fast, linewidth=3)
    # Add line at 0
    plt.axhline(0, linewidth=2, color="black", linestyle="dashed")
    x = np.arange(len(array[0, :, :].flatten()[5:]))
    # add variance as shaded area
    if var is not None:
        plt.fill_between(x, array[0, :, :].flatten()[5:] - var[0, :, :].flatten()[5:], array[0, :, :].flatten()[5:] + var[0, :, :].flatten()[5:]
                         , color=color_slow, alpha =0.25)
        plt.fill_between(x, array[1, :, :].flatten()[5:] - var[1, :, :].flatten()[5:],
                         array[1, :, :].flatten()[5:] + var[1, :, :].flatten()[5:]
                         , color=color_fast, alpha=0.25)

# Check if stimulated movements overlap with post-hoc determined slow/fast movements
for i in range(14):
    plt.figure()
    plt.subplot(2, 1, 1)
    plt.plot(stim[i, 0, :].flatten(), label="stimulated")
    plt.plot(slow[i, 0, :].flatten(), label="slow")
    plt.legend()
    plt.subplot(2, 1, 2)
    plt.plot(stim[i, 1, :].flatten(), label="stimulated")
    plt.plot(fast[i, 1, :].flatten(), label="fast")
    plt.legend()
    plt.show()


from django.conf import settings
from django.contrib.auth import get_user_model
from django.db import models

from taggit.managers import TaggableManager

User = get_user_model()

class Group(models.Model):
    title = models.CharField(max_length=200)
    slug = models.SlugField(max_length=20, unique=True)
    description = models.TextField()

    def __str__(self):
        return self.title

    class Meta:
        verbose_name = 'Community'
        verbose_name_plural = 'Communities'
        ordering = ('title',)

class Post(models.Model):
    text = models.TextField(
        'Post text',
        help_text='Enter your post text'
    )
    pub_date = models.DateTimeField(
        'Publication date',
        auto_now_add=True
    )
    author = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        related_name='posts',
        verbose_name='Author'
    )
    group = models.ForeignKey(
        Group,
        blank=True,
        null=True,
        on_delete=models.SET_NULL,
        related_name='posts',
        verbose_name='Group',
        help_text='Group to which the post belongs'
    )
    image = models.ImageField(
        'Picture',
        upload_to='posts/',
        blank=True
    )
    tags = TaggableManager(blank=True)

    def __str__(self):
        return self.text[:settings.POST_STRING_TITLE]

    class Meta:
        ordering = ('-pub_date',)
        verbose_name = 'Fast'
        verbose_name_plural = 'Posts'

class Comment(models.Model):
    post = models.ForeignKey(
        Post,
        related_name='comments',
        on_delete=models.CASCADE,
        verbose_name='A comment',
    )
    author = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        related_name='comments',
        verbose_name='Comment author'
    )
    text = models.TextField(
        'Comment text',
        help_text='New comment text'
    )
    created = models.DateTimeField(
        'Comment date',
        auto_now_add=True
    )

    def __str__(self):
        return self.text[:settings.POST_STRING_TITLE]

    class Meta:
        verbose_name = 'A comment'
        verbose_name_plural = 'Comments'
        ordering = ('-created',)

class Follow(models.Model):
    user = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        related_name='follower',
        verbose_name='User',
    )
    author = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        related_name='following',
        verbose_name='Author'
    )

    def __str__(self):
        return f'{self.user} -> {self.author.username}'

    class Meta:
        verbose_name = 'Subscription'
        verbose_name_plural = 'Subscriptions'
        unique_together = ['user', 'author']

class Like(models.Model):
    user = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        related_name='liker',
        verbose_name='User',
    )
    post = models.ForeignKey(
        Post,
        on_delete=models.CASCADE,
        related_name='liked',
        verbose_name='Fast'
    )

    def __str__(self):
        return f'{self.user} -> {self.post}'

    class Meta:
        verbose_name = 'Like'
        verbose_name_plural = 'Likes'
        unique_together = ['user', 'post']