import panel as pn
import numpy as np
import pandas as pd
import hvplot.pandas
from numba import jit
import holoviews as hv
from functools import partial

import hvplot.networkx as hvnx
import networkx as nx


@jit(nopython = True)
def gaussianC(C):

    if C != 1.0:
        Cj = np.random.uniform()
        if C<Cj:
            y = np.random.normal()
        else:
            y = 0.0
        return y
    
    else:
        
        return np.random.normal()
    return y   

#@pn.cache()
def fixed_parameters():
    
    N = 60 #0 
    R = 1e6 #1
    c = 0.2 #2
    ld = 0.1 #3
    Ng = 1e4 #4
    L = 10 #5
    K = 500 #6
    b = 5e-3 #7

    return np.array([N,R,c,ld,Ng,L,K,b])

def start_sim(parvalues,slides,button,tabs,revent):
    button.disabled = True
   
    #print(SN)
    Nsp = int(parvalues[0])
    R = slides[0].value #parvalues[1]
    c = slides[1].value #parvalues[2]
    ld = parvalues[3]
    Ngen = int(parvalues[4])
    L = int(parvalues[5])
    K = int(parvalues[6])
    b = parvalues[7]

    #print(c,R)

    TR = -1*np.ones((Nsp,L),dtype=np.int32)
    M = Mab(K);
    #print(K,L)
    G = np.zeros((Nsp, Nsp));  #Rates
    S = np.zeros((Nsp, Nsp));  #Scores
    A = np.zeros((Nsp, Nsp));  #Competition
    F = np.zeros((Nsp, Nsp));  #Foraging Efforts
    N = np.zeros(Nsp); #Populations

    #Initial Population: No
    #Minimum population allowed: Nd
    No = 1.0; Nd = 1.0;
    
    #Initial Setup:
    prm = np.array([K,L,R,ld,b]);
    so, s1, sio, fio = initial_setting(prm,M);

    N[0] = R/ld; N[1] = No;
    S[1,0] = sio; F[1,0] = fio; A[1,1] = 1.0;
    TR[0] = so; TR[1] =s1;
    
    dF = 1.0; DN = 1.0; DT = 0.2;

    SN = pd.DataFrame(columns=['SN'])
    Snumb = []
    #Stp = []
    #SN['time'] = Stp
    SN['SN'] = Snumb

    figa = (
        SN['SN']
    ).hvplot(
        title="Number of Species",
        ylabel="Number of Species",
        xlabel="Evolutionary event",
        autorange = "y",
        color="indigo",
    )
    plotns = pn.pane.HoloViews(figa, sizing_mode="stretch_both", name="PlotNS")
    tabs[0] = plotns


    
    for k in range(Ngen):
        
        while DN>1e-2:
            while dF>0.1:
                F,G,dF = foraging_strategy(N,F,A,S,b)
            N, chk, sj, DN = pop_update(N,G,ld,DT,Nd)
            if chk == True:
                N,S,F,A,TR = pop_check(N,S,F,A,TR,sj)
            dF = 1.0
        DN=1.0
        
        Fq = F
        Nq = N

        F,S,A,TR,N = speciation(M,TR,A,S,F,K,L,c,No,N);

        NA = nspecies(N)
        parvalues[8].value = NA
        parvalues[9].value = k
        Snumb.append(NA)
        #Stp.append(k)

        if k%100 == 0:
            #SN['SN'] = []
            #SN['SN'] = Snumb
            figa = (
                pd.DataFrame(Snumb,columns=['SN'])
            ).hvplot(
                title="Number of Species",
                ylabel="Number of Species",
                xlabel="Evolutionary event",
                autorange = "y",
                color="indigo",
            )
            plotns = pn.pane.HoloViews(figa, sizing_mode="stretch_both", name="PlotNS")
            tabs[0] = plotns

        if NA == Nsp-1:
            break
    
    figa = (
        pd.DataFrame(Snumb,columns=['SN'])
    ).hvplot(
        title="Number of Species",
        ylabel="Number of Species",
        xlabel="Evolutionary event",
        #xlim=(0, Ng),
        autorange = "y",
        color="indigo",
    )
    plotns = pn.pane.HoloViews(figa, sizing_mode="stretch_both", name="PlotNS")
    tabs[0] = plotns

    FwG, pos,sz,cn = get_foodweb(Fq,Nq)
    #(Fq,Nq)
    nodes = hvnx.draw_networkx_nodes(FwG, pos,with_labels=True,node_size=sz,
                                     node_color=cn)
    egx = FwG.edges()
    weights = [FwG[u][v]['edge_width'] for u,v in egx ]
    print(weights)

    edges = hvnx.draw_networkx_edges(FwG, pos, node_size=sz,edge_color='red',
                                      arrowstyle="->",arrowsize=1,alpha=0.3, edge_width=weights)
  
    figb = nodes * edges
    plotfw = pn.pane.HoloViews(figb, sizing_mode="stretch_both", name="PlotFW")
    tabs[1] = plotfw

    button.disabled = False

def Mab(K):

    A=np.ones((K,K))
    
    Iup = np.triu_indices(K);
    Idn = np.tril_indices(K);
    
    N = len(Iup[0])
    U = np.array([gaussianC(1.0) for j in range(N)])
    
    A[Iup] = U
    for i,j in zip(Idn[0],Idn[1]):
        A[i,j] = - A[j,i]
    
    np.fill_diagonal(A,0)
    return A


@jit(nopython = True)
def initial_setting(prm,M):
    
    K = prm[0]; L = prm[1]; R = prm[2];
    ld = prm[3]; b = prm[4];
    S1o = 0.0;
    bo = b/ld;
    while (S1o <= bo):
        so =  gen_string(L,K);
        s1 =  gen_string(L,K);
        S1o = Sij(so,s1,M)/L;
        f1o = 1.0;   
    
    return so, s1, S1o, f1o

@jit(nopython = True)
def gen_string(L,K):
    
    u = K*np.ones(int(L))
    l = 0

    while l<L:
        nj = np.random.randint(0,int(K))
        if nj not in u:
            u[l]=nj
            l+=1
    return u.astype(np.int32)


@jit(nopython = True)
def Sij(a,b,M):
    
    Sab = 0.0
    for ia in a:
        S = M[ia,b]
        Sab+=np.sum(S)

    if Sab > 0.0:
        return Sab
    
    else:
        return 0.0

@jit(nopython = True)
def intersect(a,b):
    
    u = np.intersect1d(a,b);
    d = len(u)/len(a)
    return d

@jit(nopython = True)
def comp_score(qij,c):
    return c+(1-c)*qij


@jit(nopython = True)
def foraging_strategy(N,F,A,S,b):

    Fn = np.zeros(F.shape)
    G = np.zeros(F.shape)
    nk = np.where(N>0)[0]
    jn = np.max(nk)
    
    for i in nk:
        if(i>0):
            Sq = S[i,:]
            Fq = F[i,:]
            pk = np.where(Sq>0)[0]

            for j in pk:
                Sk = np.where(S[:,j]>0)[0]

                sn = 0.0;
                for k in Sk:
                    sn+=A[k,i]*S[k,j]*F[k,j]*N[k]
                G[i,j]=Sq[j]*Fq[j]*N[j]/(b*N[j]+sn)
    
    for i in nk:
        if (i>0):
            SGij = np.sum(G[i,:])
            Fn[i,:] = G[i,:]/SGij
        jn = np.where(Fn[i,:]>0.0)
        for j in jn[0]:
            if Fn[i,j]<1e-6:
                Fn[i,j] =  1e-6
                    

    DF = np.max(np.abs((F-Fn).ravel()))

    return Fn, G, DF


@jit(nopython = True)
def pop_update(N,G,ld,DT,Nd):

    rem = False;
    
    Np = np.zeros(N.shape)
    la = (1 - DT)*N;
    lb = ld*DT*np.sum(G,axis=1)*N
    V = np.zeros(N.shape)

    nk = np.where(N>0)[0]

    for j in nk:
        if j>0:
            V[j]=np.sum(G[:,j]*N)
    
    Np = la + lb - DT*V
    Np[0] = N[0]

    U = Np[(Np!=0) & (Np<Nd)]
    if len(U)>0:
        sj =  np.where(((Np!=0) & (Np<Nd)))[0]
        rem = True
        Np[sj] = 0.0

    DN = np.max(np.abs(N-Np))

    return Np,rem, sj, DN

@jit(nopython = True)
def pop_check(N,S,F,A,TR,jn):
    Ln = len(TR[0,:])
    for jd in jn:
        N[jd] = 0.0
        S[jd,:] = 0.0
        S[:,jd] = 0.0
        F[jd,:] = 0.0
        F[:,jd] = 0.0
        A[jd,:] = 0.0
        A[:,jd] = 0.0
        TR[jd,:] = -1*np.ones(Ln,dtype=np.int32)

    return N,S,F,A,TR  


@jit(nopython = True)
def speciation(M,TR,A,S,F,K,L,c,No,N):
    Tn = TR[:,0]

    nsj = np.min(np.where(Tn==-1)[0]) #Allocated new species 
    Na = np.where(Tn>-1)[0]
    Nb = Na[Na!=0] #excludes external species
    ns = np.random.choice(Nb) #parent species
    maxj = np.max(Nb) #index of largest non zero species

    snew = False; NFT = False;
    new_traits = TR[ns,:].copy()

    while(snew==False):
        while(NFT == False):
            nf = np.random.randint(0,K)
            if nf not in new_traits:
                jn = np.random.randint(0,L)
                new_traits[jn] = nf
                NFT = True
        n=0
        for q in Na:
            a = TR[q]
            dij = np.intersect1d(a,new_traits)
            lj = len(dij)
            if (lj == L):
                n+=1
        
        if n==0:
            snew = True
        else:
            NFT = False

    ##Initialising scores and pop
    TR[nsj] = new_traits
    N[nsj] = No
    N[ns] -= 1 
    A[nsj][nsj] = 1.0

    ia = TR[nsj]
    
    for q in Na:
        ja = TR[q]

        A[nsj][q] = comp_score(intersect(ia,ja),c)
        A[q][nsj] = A[nsj][q]

        if q == 0:
       
            S[nsj][0] = Sij(ia,ja,M)
            if S[nsj][0] > 0.0:
               
                if F[ns][0] > 0.0:
                    F[nsj][0] = F[ns][0]
                else:
                    F[nsj][0] = np.random.uniform(1e-6,1)
        
        if q>0:
            S[nsj][q] = Sij(ia,ja,M)
            S[q][nsj] = Sij(ja,ia,M)

            if S[nsj][q] > 0.0:
                if F[ns][q] > 0.0:
                    F[nsj][q] = F[ns][q]
                else:
                    F[nsj][q] = np.random.uniform(1e-6,1)
            
            if S[q][nsj]>0.0:
                if F[q][ns] > 0.0:
                    F[q][nsj] = F[q][ns]
                else:
                    F[q][nsj] = np.random.uniform(1e-6,1)
    
    return F,S,A,TR,N


@jit(nopython = True)
def nspecies(N):

    Na = np.where(N>0)[0]
    #ns =np.max(Na)
    NA= len(Na)
    
    return NA-1

def get_foodweb(F,N):
    
    F[F<=0.01] = 0.0
    #Positive Population indices
    jall = np.where(N>0)[0]
    #All positive pop indices without the external species 
    jy = jall[jall>0]     
    #Basal Species indices
    jo = np.where(F[:,0]>0.0)[0] 
    tlvs = []
    tlvs.append(jo)
    spn = False
    M=0
    #Levels
    while spn == False:
        jx = np.setdiff1d(jy,jo) #Speciesn in jy but not in jo
        ln=[]
        for q in jx:
            Un = np.where(F[q,:]>0.0)[0]
            for p in jo:
                if p in Un:
                    ln.append(q)
                    break
        ln=list(set(ln))
        tlvs.append(ln)
        un = list(jo)+ln
        jo = np.asarray(un)
        M+=1
        
        if M ==5:
            spn = True
    
    DN = []
    DR = []
    DRL = []
    for q in tlvs:
        if len(q)>0:
            Plog = np.log(1+N[q])
            Nlv = len(Plog)
            DN.append(Plog)
            DR.append(np.max(Plog))
            DRL.append(Nlv)
    
    Dy = np.max(DR)
    
    Xn = []; Yn = [];
    for q in range(len(DN)):
        X = [(6*j+1)*DR[q]-DR[q] for j in range(DRL[q])]
        Y = [(q+1)*8*Dy for j in range(DRL[q])]
        Xn.append(X);Yn.append(Y)

    Lmx = []    
    for q in Xn:    
        Lmx.append(q[len(q)-1])
        
    lj = np.max(Lmx)

    for q in range(len(Xn)):
        if len(Xn[q]) > 1:
            Xn[q] = [lj/Lmx[q]*k for k in Xn[q]]
    
    FW = pd.DataFrame(columns=['n','X','Y','R'])
    sn = [];xn = []; yn = []; rn = []
    for q in range(len(tlvs)):
        for p in range(len(tlvs[q])):
            sn.append(tlvs[q][p])
            xn.append(Xn[q][p])
            yn.append(Yn[q][p])
            rn.append(DN[q][p])

    FW['n'] = sn; FW['X'] = xn; FW['Y'] = yn; FW['R'] = rn
    
    tlevs = []

    for q in tlvs:
        if len(q)>0:
            tlevs.append(q)
    
    
    G = nx.DiGraph()
    pos = {}
    n=0
    sizes = []
    colormap = []
    edgesw = []
    
    for q in tlevs:
        if n == 0:
            basal  = FW['n'].isin(q)
            FBAS = FW[basal]
            RY = np.max(list(FBAS['R']))
            
            for i in q:
                
                nprop = FBAS.loc[FBAS['n']==i]
                rn = nprop['R'].values[0]
                xn = nprop['X'].values[0]
                yn = nprop['Y'].values[0]
                
                a = 'er'+str(i)
                b = str(i)
                G.add_node(a)
                sizes.append(100.0)
                pos[a] = (xn,0.0)
                colormap.append('yellow')
                G.add_node(b)
                sizes.append(30*rn)
                colormap.append('blue')    
                pos[b] = (xn,yn)
                G.add_edge(a,b,edge_width = 5*F[i,0])
            n+=1
        else:

            TRL = FW['n'].isin(q)
            FRL = FW[TRL]

            for i in q:
                a = str(i)
                nprop = FRL.loc[FRL['n']==i]
                rn = nprop['R'].values[0]
                xn = nprop['X'].values[0]
                yn = nprop['Y'].values[0]
                G.add_node(a)
                sizes.append(30*rn)
                colormap.append('blue')    
                pos[a] = (xn, yn)
        
        for q in tlevs:
            for i in q:
                jn = np.where(F[i,:]>0)[0]
                if len(jn)>0:
                    for j in jn :
                        if j!=0:
                            a=str(j);b=str(i)
                            G.add_edge(a,b,edge_width=5*F[i,j])

    return G,pos,sizes,colormap