import numpy as np
from scipy.integrate import odeint
from tqdm.notebook import tqdm
from numpy import tan,arctan,exp,sqrt,arcsin,sin,pi
import matplotlib.pyplot as plt
import random
import time
import itertools
from matplotlib.patches import Rectangle
def Inoue_Nagayoshi(d,totaln,x0=0,y0=0,s=50,A=50,A_B=10,theta=-0.5,nin=1,a=4,b=3.995,eps=.002,init_model='rand',N_freeze_limit=8,analyze_mode=False,pbplt=0,update=0):
    def activation(xvh,yvh):
        return np.where(abs(xvh-yvh) <eps, 1, 0)
    def f(xv):
        return a*xv*(1-xv)
    def g(yv):
        return b*yv*(1-yv)
    def eom(xv,yv,I):
        return ((f(xv)*(1+I)+I*g(yv))/(2*I+1),(g(yv)*(1+I)+I*f(xv))/(2*I+1))
    def delta(ii,jj):
        return int(ii==jj)
    def init():
        if init_model == 'rand':
            xj=(np.random.random((N,N)))
            yj=xj+(np.random.random((N,N))-0.5)/100
            return (xj,yj,activation(xj,yj))
        if init_model == 'given':
            return (x0,y0,activation(x0,y0))
    def Energy(uv,wv,wv1):
        return [-np.einsum('ijkl,ij,kl', wv1, uv, uv) / 2,-np.einsum('ijkl,ij,kl', wv, uv, uv) / 2 - s * np.sum(uv)+N*A]
    
    B=A/A_B
    N=len(d)
    xs,ys,u=init()
    I=np.zeros((N,N))
    ar=np.arange(N)
    Es2=np.zeros(totaln)
    w=np.zeros((N,N,N,N))
    w1=np.zeros((N,N,N,N))
    for i, k, j, l in itertools.product(ar, repeat=4):
        w1[i,k,j,l]=-A*(delta(i,j)*(1-delta(k,l))+delta(k,l)*(1-delta(i,j)))
        w[i,k,j,l]=w1[i,k,j,l]-B*d[i,j]*(delta(l,(k+1)%N)+delta(l,(k-1)%N))
    N_valid=0
    Step=totaln
    Emin = Energy(u,w,w1)[1]
    ubest=u
    u_freeze = u
    N_freeze = 0
    Ns=np.arange(totaln)
    tt=time.time()
    Time=0
    
    for n in Ns[1:]: 
        if pbplt!=0:
            if ((n+1)/totaln*100%10<=0.001):
                pbplt.clear()
                pbplt.add_patch(Rectangle((0, 0),n/totaln, 1))
                update.flush_events()
                update.draw()
        I=np.einsum('ijkl,kl->ij', w, u)+s-theta
        I=np.where(I>0,I,0)
        for inter in range(nin):
            for i, j in itertools.product(ar, repeat=2):
                xs[i,j],ys[i,j]=eom(xs[i,j],ys[i,j],I[i,j])
        u = activation(xs,ys)
        if np.array_equal(u,u_freeze):
            N_freeze+=1
            if N_freeze_limit>N_freeze_limit:
                xs,ys,u=init()
                N_freeze_limit=0
        else:
            E1,E2 = Energy(u,w,w1)
            if E1==0:
                N_valid+=1
            if E2<Emin:
                Emin=E2
                ubest=u
                Step = n
                Time = time.time()-tt
        if analyze_mode:
            Es2[n] = E2
        u_freeze=u
    if analyze_mode:
        plt.plot(Ns[1:],Es2[1:])
    dic = {'Best': ubest,
      'Energy': Emin,
      'Step': Step,
      'Time': Time,
      'Valid': N_valid,
      'Method': 'Inoue_Nagayoshi'}
    return dic
def random_cities(N):
    cities=np.array([[0,0]])
    for ii in range(N):
        cities=np.append(cities,[np.random.random(2)],axis=0)
    cities=np.delete(cities,0,axis=0)
    d=np.zeros((N,N))
    for i in range(N):
        for j in range(N):
            x1,y1=cities[i]
            x2,y2=cities[j]
            d[i,j]=sqrt((x1-x2)**2+(y1-y2)**2)
    return (d,cities)
def plot_road(ud,cities,axs = plt,labels=''):
    cold = np.argmax(ud[:,0])
    axs.plot(cities[:,0],cities[:,1],'o',markersize=10)
    for j in range(len(cities)-1):
        cnew = np.argmax(ud[:,j+1])
        x1,y1=cities[cold]
        x2,y2=cities[cnew]
        axs.plot([x1,x2],[y1,y2])
        cold=cnew
    cnew = np.argmax(ud[:,0])
    x1,y1=cities[cold]
    x2,y2=cities[cnew]
    axs.plot([x1,x2],[y1,y2])
    # axs.legend()
    # axs.show()
def analyze(nni,ii,fld):
    print(fld)
    print('Number of cities: %i'%num_cities)
    print("Number of seeds: %i"%seeds)
    print('Max steps: %i'%num_steps_1)
    for st,t,nin in zip(steps[nni,ii],times[nni,ii],ninm):
        print('Nin = %i: %f steps, %f times'%(nin,np.mean(st),np.mean(t)))
    plt.plot(ninm,np.mean(steps[nni,ii],axis=1)/max(np.mean(steps[nni,ii],axis=1)),label = 'Steps')
    plt.plot(ninm,np.mean(times[nni,ii],axis=1)/max(np.mean(times[nni,ii],axis=1)),label = 'Times')
    if min(np.mean(energies[0,0],axis=1))!=max(np.mean(energies[0,0],axis=1)):
        plt.plot(ninm,(np.mean(energies[nni,ii],axis=1)-min(np.mean(energies[nni,ii],axis=1)))/(max(np.mean(energies[nni,ii],axis=1))-min(np.mean(energies[nni,ii],axis=1))),label = 'Energies')
    
    plt.ylabel('Normalized number of steps/time:')
    plt.xlabel('Nin parameter')
    plt.legend()
    plt.ylim(0,1.1)
    # plt.yscale('log')
    plt.savefig(fld+str(num_cities)+'.png')
    plt.show()
def analyze2(ii,jj,jjj):
    print('Number of cities: %i'%num_cities)
    print("Number of seeds: %i"%seeds)
    print("Number of nin: %i"%jjj)
    print('Max steps: %i'%num_steps_1)
    print('        ',NNS)
    for seed in range(seeds):
            print('Seed %i:  '%seed,end='') 
            for nnn in np.arange(len(NNS)):
                print('%i steps'%int(steps[nnn,ii,jj,seed]),end='  ')
            print('')
    print('        ',NNS)
    for seed in range(seeds):
            print('Seed %i:  '%seed,end='')
            for nnn in np.arange(len(NNS)):
                print('%.2f sec'%times[nnn,ii,jj,seed],end='  ')
            print('')
    print('        ',NNS)
    for seed in range(seeds):
            print('Seed %i:  '%seed,end='')
            for nnn in np.arange(len(NNS)):
                print('%.2f     '%energies[nnn,ii,jj,seed],end='  ')
            print('')
    for nni,NN in enumerate(NNS):
        for st,t,nin in zip(steps[nni,ii],times[nni,ii],ninm):
            print('%s: %f steps, %f times'%(NN,np.mean(st),np.mean(t)))
def rand_init(N,model='rand',T=0):
    if model=='rand':
        xj=np.random.random((N,N))-0.5
    if model=='sym':
        if T==0:
            T=1/50000
        # xj  = T*np.log[-(u/(-1 + u))]
        u00 = 1/N
        u=((np.random.random((N,N))-0.5)*0.2+1)*u00
        xj = T*np.log(-(u/(-1 + u)))
    if model=='duf':
        u00 = 1/N
        xj=((np.random.random((N,N))-0.5)*0.2+1)*u00
    return (xj)


def SinMap(d,totaln,x0=0,s=50,A=50,A_B=10,theta=0,nin=1,eta=100,T=0,init_model='rand',analyze_mode=False,pbplt=0,update=0):
    def activation(xvh):
        if T==0:
            return np.where(xvh >= 0, 1, 0)
        else:
            return 1/(1+np.exp(-xvh/T))
    def g(xv,I):
        e = 0.25/(1+eta*I**2)
        if xv<0:
            if I>0:
                e=0.5-e
            return 0.5*sin(2*xv*(pi+arcsin(2*e)))
        else:
            if I<0:
                e=0.5-e
            return 0.5*sin(2*xv*(pi+arcsin(2*e)))
    def delta(ii,jj):
        return int(ii==jj)
    def init():
        if init_model == 'rand':
            xj=np.random.random((N,N))-0.5
            return (xj,activation(xj))
        if init_model == 'given':
            return (x0,activation(x0))
    def Energy(uv,wv,wv1):
        return [-np.einsum('ijkl,ij,kl', wv1, uv, uv) / 2,-np.einsum('ijkl,ij,kl', wv, uv, uv) / 2 - s * np.sum(uv)+N*A]
    tt=time.time()
    B=A/A_B
    N=len(d)
    xs,u=init()
    I=np.zeros((N,N))
    ar=np.arange(N)
    Es2=np.zeros(totaln)
    w=np.zeros((N,N,N,N))
    w1=np.zeros((N,N,N,N))
    for i, k, j, l in itertools.product(ar, repeat=4):
        w1[i,k,j,l]=-A*(delta(i,j)*(1-delta(k,l))+delta(k,l)*(1-delta(i,j)))
        w[i,k,j,l]=w1[i,k,j,l]-B*d[i,j]*(delta(l,(k+1)%N)+delta(l,(k-1)%N))
    Step=totaln
    E1,E2 = Energy(u,w,w1)
    Emin = E2
    ubest=u
    N_valid=0
    Ns=np.arange(totaln)
    u_freeze=u
    Es2=np.zeros(totaln)
    Time=0
    tt=time.time()
    
    for n in Ns[1:]: 
        if pbplt!=0:
            if ((n+1)/totaln*100%10<=0.001):
                pbplt.clear()
                pbplt.add_patch(Rectangle((0, 0),n/totaln, 1))
                update.flush_events()
                update.draw()
            
        I=np.einsum('ijkl,kl->ij', w, u)+s+theta   
        vec_g = np.vectorize(g)
        for it in range(nin):
            xs = vec_g(xs, I)
            u = activation(xs)

        if np.array_equal(u,u_freeze) ^ True:
            E1,E2 = Energy(u,w,w1)
            if E1==0:
                N_valid+=1
            if E2<Emin:
                Emin=E2
                ubest=u
                Step = n
                Time = time.time()-tt
            u_freeze=u
        if analyze_mode:
            Es2[n] = E2
    if analyze_mode:
        plt.plot(Ns[1:],Es2[1:])
    dic = {'Best': ubest,
      'Energy': Emin,
      'Step': Step,
      'Time': Time,
      'Valid': N_valid,
      'Method': 'Sinmap'}
    return dic

def Duffing(d,totaln,x0=0,y0=0,s=50,A=50,A_B=10,theta=0,nin=4,eta=100,T=0,init_model='rand',Lambda=0.25,ddt=0.01,internal=60,analyze_mode=False,pbplt=0,update=0): 
    def activation(xvh):
        if T==0:
            return np.where(xvh >= 0, 1, 0)
        else:
            return 1/(1+exp(-xvh/T))
    def delta(ii,jj):
        return int(ii==jj)
        
    def init():
        if init_model == 'rand':
            xj=np.random.random((N,N))-0.5
            yj=np.random.random((N,N))-0.5
            return (xj,yj,activation(xj))
        if init_model == 'given':
            return (x0,y0,activation(x0))
        if init_model == 'eye':
            xj=torch.eye(N)-0.5
            yj=torch.eye(N)-0.5
            return (xj,yj,activation(xj))
    def model_disipative(u, t,ee,eb):
        return (u[1], ee-u[1]*alpha + beta*u[0] -gamma*(u[0])**3 + f*np.cos(omega*t))
    def xn(Tend,e,dt,xold,yold):
        x0=xold
        y0=yold
        if Tend==1:
            tT = np.linspace(0,(internal+Tend)*dt,2)
            z1 = odeint(model_disipative, [x0, y0], tT,args=(e,e))
        else:
            t0 = (Tend-1+internal)*dt
            for ty in range(int(dt/ddtt)):
                der = model_disipative([x0,y0],t0,e,e)
                x0 = x0 + der[0]*ddtt
                y0 = y0 +der[1]*ddtt
                t0 = t0 + ddtt
        return (x0,y0)
    def Energy(uv,wv,wv1):
        return [-np.einsum('ijkl,ij,kl', wv1, uv, uv) / 2,-np.einsum('ijkl,ij,kl', wv, uv, uv) / 2 - s * np.sum(uv)+N*A]
    
    f=1
    alpha=1
    beta=10
    gamma=100
    omega=3.5
    B=A/A_B
    N=len(d)
    xs,ys,u=init()
    I=np.zeros((N,N))
    ar=np.arange(N)
    Es2=np.zeros(totaln)
    w=np.zeros((N,N,N,N))
    w1=np.zeros((N,N,N,N))
    for i, k, j, l in itertools.product(ar, repeat=4):
        w1[i,k,j,l]=-A*(delta(i,j)*(1-delta(k,l))+delta(k,l)*(1-delta(i,j)))
        w[i,k,j,l]=w1[i,k,j,l]-B*d[i,j]*(delta(l,(k+1)%N)+delta(l,(k-1)%N))
    N_valid=0
    u_freeze=u
    Step=totaln
    E1,E2 = Energy(u,w,w1)
    Emin = E2
    ubest=u
    Ns=np.arange(totaln)
    tt=time.time()
    Time=0
    for n in Ns[1:]:
        if pbplt!=0:
            if ((n+1)/totaln*100%10<=0.001):
                pbplt.clear()
                pbplt.add_patch(Rectangle((0, 0),n/totaln, 1))
                update.flush_events()
                update.draw()
        I=np.einsum('ijkl,kl->ij', w, u)+s+theta
        eps=Lambda*arctan(I)
        vec_g = np.vectorize(xn)
        xs,ys = vec_g(n,eps,ddt,xs,ys)
        u = activation(xs)
        
        if np.array_equal(u,u_freeze) ^ True:
            E1,E2 = Energy(u,w,w1)
            if E1==0:
                N_valid+=1
            if E2<Emin:
                Emin=E2
                ubest=u
                Step = n
                Time = time.time()-tt
            u_freeze=u
        if analyze_mode:
            Es2[n] = E2
    if analyze_mode:
        plt.plot(Ns[1:],Es2[1:])
    dic = {'Best': ubest,
      'Energy': Emin,
      'Step': Step,
      'Time': Time,
      'Valid': N_valid,
      'Method': 'Duffing'}
    return dic