import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
import math
import time as time

# special points on the surface

# minima
min3pos = -0.558224, 1.44173
min1pos = 0.623499, 0.0280378
min2pos = -0.0500108, 0.466694
mainminpos = min1pos

# saddle points
saddle_pos = np.array([0.212487,0.292988])
saddle_pos2 = np.array([-0.822002, 0.624313])
mainsaddlepos = saddle_pos


#%%Initiatlization

time_begin = time.clock()

coefficients = np.array([[-200, -1, 0, -10, 1, 0],
                         [-100, -1, 0, -10, 0, 0.5],
                         [-170, -6.5, 11, -6.5, -0.5, 1.5],
                         [15, 0.7, 0.6, 0.7, -1, 1]])

pointnum = 100                                                                   #how fine shoutld the grid of the contour plot be
V = np.zeros((pointnum,pointnum))

x_min2 = -1.5                                                                    #max and min x and y values for plot
x_max2 = 0.2
y_min2 = 0.2
y_max2 = 1.8

x_min = -0.5                                                                     #max and min x and y values for plot
x_max = 1.2
y_min = -0.3
y_max = 0.8

buffer = 0.005                                                                     #time particle can penetrate into the other valley - to make sure it stays there

x = np.linspace(x_min,x_max,pointnum)
y = np.linspace(y_min,y_max,pointnum)
X, Y = np.meshgrid(x, y)                                                         #make X,Y grid

def energy(x,y):
    a = sum(coeff[0]*np.exp(coeff[1]*(x-coeff[4])**2 + coeff[2]*(x-coeff[4])*(y-coeff[5])
    + coeff[3]*(y-coeff[5])**2) for coeff in coefficients)                       #calculate Müller Brown PES
    return a

def force(x,y):
    Fx = -sum((coeff[0]*np.exp(coeff[1]*(x-coeff[4])**2 + coeff[2]*(x-coeff[4])*(y-coeff[5]) +
               coeff[3]*(y-coeff[5])**2))*(coeff[1]*(x-coeff[4])*2 + coeff[2]*(y-coeff[5])) for coeff in coefficients)
    Fy = -sum((coeff[0]*np.exp(coeff[1]*(x-coeff[4])**2 + coeff[2]*(x-coeff[4])*(y-coeff[5]) +
               coeff[3]*(y-coeff[5])**2))*(coeff[3]*(y-coeff[5])*2 + coeff[2]*(x-coeff[4])) for coeff in coefficients)
    return Fx, Fy

def Hessian(x,y,delta):
    Hess = np.array([[(force(x,y)[0]-force(x+delta,y)[0])/delta, (force(x,y)[1]-force(x+delta,y)[1])/delta],
                     [(force(x,y)[0]-force(x,y+delta)[0])/delta, (force(x,y)[1]-force(x,y+delta)[1])/delta]])
    return Hess

def calcZ(mainminpos,min2pos,betha,deltaEmins):
    Hmin1 = Hessian(mainminpos[0],mainminpos[1],1e-7)
    Hmin2 = Hessian(min2pos[0],min2pos[1],1e-7)
    curvatures1, curvatures2 = np.linalg.eig(Hmin1)[0],np.linalg.eig(Hmin2)[0]
    Z = 2*np.pi/betha/(np.sqrt(curvatures1[0]*curvatures1[1])+np.sqrt(curvatures2[0]*curvatures2[1])*np.exp(-betha*deltaEmins))
    Z = Z*np.pi*2/betha/m
    return Z

def calZalternative(pointnum,V_backup,x_min,x_max,y_min,y_max):
    Zalt = np.sum(np.exp(-betha*(V_backup-energy(mainminpos[0],mainminpos[1]))))*(x_min-x_max)*(y_min-y_max)/(pointnum**2)
    Zalt = Zalt * np.pi*2/betha/m
    return Zalt

def calcZinResPhaseSpace(pointnum,V_backup,x_min,x_max,y_min,y_max):
    V_backup[V_backup<V_backup.min()+deltaE] = V_backup.min()+deltaE
    Z = np.sum(np.exp(-betha*(V_backup-energy(mainminpos[0],mainminpos[1]))))*(x_min-x_max)*(y_min-y_max)/(pointnum**2)
    Z = Z*np.pi*2/betha/m
    return Z

def computevec(mainminpos, mainsaddlepos):
    H = Hessian(mainsaddlepos[0],mainsaddlepos[1], 10e-7)
    Eigs = np.linalg.eig(H)
    whichvec = int(np.where(Eigs[0]<0)[0])                                            # which vector has negative eigenvalue
    vec = Eigs[1][:,whichvec]
    vec = vec*np.sign(np.dot(vec,mainminpos - mainsaddlepos))                         # gets multiplied with -1 if it is looking away from the mainmin
    return vec

def calcP(deltaE,betha,Z,mainsaddlepos):
    H = Hessian(mainsaddlepos[0],mainsaddlepos[1], 10e-7)
    V = np.linalg.eig(H)[0]
    P = 2*np.pi*np.exp(-betha*deltaE)/(betha**2*m**(3/2)*np.sqrt(float(V[V>0])))/Z
    return P

def checkside(vec, xy_data,mainsaddlepos, buffer, delta_t):
    buffer_steps = int(buffer/delta_t)
    crossboarder = 0
    dist_sepplane = np.matmul(np.transpose(xy_data) - mainsaddlepos.reshape(1,2) , vec.reshape(2,1))
    dist_sepplane[dist_sepplane < 0] = -1
    dist_sepplane[dist_sepplane > 0] = 1

    change = dist_sepplane - np.row_stack((dist_sepplane[0],dist_sepplane[0:np.size(dist_sepplane)-1]))
    if np.any(change != 0):
        if np.sum(dist_sepplane[np.where(change != 0)[0][0]:np.where(change != 0)[0][0] + buffer_steps]) == -buffer_steps:
            crossboarder = 1
    return crossboarder

def checkside_with_comeback(vec, xy_data):
    crossboarder = 0
    dist_sepplane = np.matmul(np.transpose(xy_data) - mainsaddlepos.reshape(1,2) , vec.reshape(2,1))
    dist_sepplane[dist_sepplane < 0] = 0                                                                       # towards sidemin
    dist_sepplane[dist_sepplane > 0] = 1                                                                       # towards mainmin

    change = dist_sepplane - np.row_stack((dist_sepplane[0],dist_sepplane[0:np.size(dist_sepplane)-1]))
    changesum = np.sum(change[change != 0])
    if changesum != 0:
        if np.sum(dist_sepplane[np.where(change != 0)[0][0]:np.where(change != 0)[0][0] + 5]) == 0 or np.sum(dist_sepplane[np.where(change != 0)[0][0]:np.where(change != 0)[0][0] + 5]) == 5:
            crossboarder = 1
    return crossboarder


def calcMat(betha,pointnum,xmin,xmax,ymin,ymax,deltaE):
    x = np.linspace(xmin,xmax,pointnum*3)
    y = np.linspace(ymin,ymax,pointnum*3)
    X, Y = np.meshgrid(x, y)
    Vacc = np.exp(-betha*energy(X,Y))
    Emin = energy(mainminpos[0],mainminpos[1])

    Vacc[Vacc>np.exp(-betha*(Emin+deltaE))] = np.exp(-betha*(Emin + deltaE))

    VaccCUM = np.cumsum(Vacc)
    VaccCUMSUM = VaccCUM.reshape(pointnum*3,pointnum*3)
    VaccCUMSUM_norm = VaccCUMSUM/VaccCUMSUM.max()
    return VaccCUMSUM_norm, VaccCUM


def findStartv(betha,m):
    v_abs = np.sqrt(-2/(betha*m)*np.log(1-np.random.uniform(0,1)))
    rand = np.random.uniform(0,2*np.pi)
    randvec = np.array([np.cos(rand),np.sin(rand)])
    v_maxboltz = v_abs * randvec
    return v_maxboltz


def findStartv_restricted_phase_space(betha,m,startE,deltaE):
    E_kinmin = float(np.array([deltaE-startE,0]).max())
    v1 = np.sqrt(2*(E_kinmin)/m)
    v_abs = np.sqrt(v1**2 - 2*np.log(1- np.random.uniform(0,1))/(betha*m))
    rand = np.random.uniform(0,2*np.pi)
    randvec = np.array([np.cos(rand),np.sin(rand)])
    v_maxboltz = v_abs * randvec
    return v_maxboltz


def findStartPos(VaccCUMSUM_norm,xmin,xmax,ymin,ymax,pointnum):
    r = np.random.uniform(0,1)
    Vdiff_row = np.abs(VaccCUMSUM_norm[:,0] - r)
    nearV = Vdiff_row.min()
    ygrid = int(np.where(Vdiff_row  < nearV + 1e-18)[0][0])

    Vdiff_zeile = np.abs(VaccCUMSUM_norm[ygrid-1:ygrid+1,:] - r)
    nearestV = np.abs(Vdiff_zeile.min())
    xgrid = int(np.where(Vdiff_zeile < nearestV + 1e-20)[1][0])
    ygrid = ygrid - 1 + np.where(Vdiff_zeile < nearestV + 1e-15)[0][0]
    y =  ymin + ygrid*(ymax-ymin)/(pointnum*3)
    x =  xmin + xgrid*(xmax-xmin)/(pointnum*3)
    return x,y


def findMinEnergyPath(m,mainsaddlepos):
    x = np.zeros(1000+1)
    y = np.zeros(1000+1)
    delta = 0.0001
    x[0],y[0] = mainsaddlepos + (mainminpos - mainsaddlepos)*delta
    for i in np.linspace(1,1000,1000):
        i = int(i)
        Fx = -sum((coeff[0]*np.exp(coeff[1]*(x[i-1]-coeff[4])**2 + coeff[2]*(x[i-1]-coeff[4])*(y[i-1]-coeff[5]) +
               coeff[3]*(y[i-1]-coeff[5])**2))*(coeff[1]*(x[i-1]-coeff[4])*2 + coeff[2]*(y[i-1]-coeff[5])) for coeff in coefficients)
        Fy = -sum((coeff[0]*np.exp(coeff[1]*(x[i-1]-coeff[4])**2 + coeff[2]*(x[i-1]-coeff[4])*(y[i-1]-coeff[5]) +
               coeff[3]*(y[i-1]-coeff[5])**2))*(coeff[3]*(y[i-1]-coeff[5])*2 + coeff[2]*(x[i-1]-coeff[4])) for coeff in coefficients)
        x[i], y[i] = x[i-1] + Fx*delta, y[i-1] + Fy*delta

        if x[i] - x[i-1] < delta*1e-2 and y[i] - y[i-1] < delta*1e-2 and i>20:
            break

    deltax, deltay = x[i] - x[i-1], y[i] - y[i-1]
    deltax, deltay = deltax/(np.sqrt(deltax**2+deltay**2)), deltay/(np.sqrt(deltax**2+deltay**2))
    deltaxy = np.array([deltax, deltay]).reshape(2,1)
    return deltaxy, x[i], y[i]


def calcArrheniusTransRate(deltaxy, m, deltaE, x_last, y_last, betha):
    H = Hessian(x_last,y_last,1e-7)
    gradchange = np.matmul(H,deltaxy)
    curv = np.matmul(gradchange.reshape(1,2),deltaxy)
    w = np.sqrt(curv/m)
    k_arrhenius = w/(2*np.pi)*np.exp(-betha*deltaE)
    return k_arrhenius


V_backup = energy(X,Y)
V = np.array(V_backup)                                                           #calculate the PES by calling energy function on the X,Y grid
V[V>200] = float('nan')                                                          #set very big values of PES to NaN for estetical reasons


#%%FIGURE

fig = plt.figure()
#ax = fig.add_subplot(111, projection='3d') #fig.projection = '3d'               #uncomment if you want 3D plot
CS = plt.contour(X, Y, V, pointnum, linewidths=0.5, colors='k')                  #contour plot with xx contour lines
CS = plt.contourf(X, Y, V, pointnum)                                             #adds colors to the contour lines
plt.colorbar()                                                                   #draw colorbar



#%%IMPLEMENT EQUATIONS OF MOTION

deltaE = -energy(mainminpos[0],mainminpos[1]) + energy(mainsaddlepos[0],mainsaddlepos[1])     #energy difference between saddle point and minimum (small two minimas and their saddle)
deltaEmins = energy(min2pos[0],min2pos[1]) - energy(mainminpos[0],mainminpos[1])              #energy diff between the two minima
betha = 1e-1*1.6
m = 1
delta_t = 0.001                                                                               #one increment of time

l = -1
times = np.array([0.1]) + buffer                                                  #times one particle can propagate till thermal equilibration is reached --> then it gets new position and velocities
predtrans = np.zeros(int(np.size(times)))
actualtrans = np.zeros(int(np.size(times)))

VaccCUMSUM_norm = calcMat(betha,pointnum,x_min,x_max,y_min,y_max,deltaE)[0]



traj_num_max = 20000

E_diff = np.zeros(traj_num_max+1)                                                 # will be the difference of starting energy and final energy after all runs
traj_num = -1
pottrans_traj_num = 0
trans_traj_num = np.zeros(np.size(times))
count = 0

count_occurrances = np.zeros([pointnum,pointnum])

x,y = mainminpos[0],mainminpos[1]
vx,vy = 5,8
gamma = 30
v = np.zeros(int((traj_num_max*((times[0] + buffer+0.001)/delta_t ))))
sigma = np.sqrt(2*gamma/(betha*delta_t))

while traj_num<traj_num_max:                                                      #traj_num < number of trajectories we are gonna sample
    traj_num+=1
    t = 0
    j = -1

    U = (energy(x,y)-V_backup.min())

    T = m/2*(vx**2 +vy**2)
    #print(E)
    vx, vy = vx + force(x,y)[0]/m*delta_t/2, vy + force(x,y)[1]/m*delta_t         #make a half step at beginning (leap frog)

    E = U+T                                                                       #calculate total Energy
    xy_data = np.zeros((2,int(times[0]/delta_t)+1))                               #initialize xy_array


    while t<times[0]+delta_t:
        j+=1
        count+=1
        x,y = x+vx*delta_t, y+vy*delta_t
                                                                                  #equations of motion
        Fx_cons, Fy_cons = force(x,y)
        Fx_fric, Fy_fric = -gamma*vx,-gamma*vy
        Fx_rand, Fy_rand = np.random.normal(0,sigma), np.random.normal(0,sigma)
        Fx_net,Fy_net = Fx_rand+Fx_cons+Fx_fric, Fy_rand+Fy_cons+Fy_fric

        vx, vy = vx+Fx_net/m*delta_t, vy+Fy_net/m*delta_t

        v[count] = vx

        xy_data[:,j] = x,y
        t = t+delta_t
        if x<x_max and y<y_max:
            x_grid = int((x-x_min)*pointnum/(x_max-x_min))
            y_grid = int((y-y_min)*pointnum/(y_max-y_min))
            count_occurrances[x_grid,y_grid] = count_occurrances[x_grid,y_grid] + 1

    U2 = (energy(x+vx*delta_t/2,y+vy*delta_t/2)-V_backup.min())               #calculate pot and kin energy after the sweep
    T2 = m/2*(vx**2 +vy**2)
    E2 = U2+T2
    E_diff[traj_num] = np.abs(E-E2)                                           #calc numerical error in energy

    if checkside(computevec(mainminpos,mainsaddlepos) , xy_data[:,0:int(times[0]/delta_t)],mainsaddlepos,buffer,delta_t) == 1:
        np.savetxt('x_data.dat',xy_data[0,:])                                #save x and y Data for scatter plot of animation
        np.savetxt('y_data.dat',xy_data[1,:])
        print(1)
        trans_traj_num[0] = trans_traj_num[0]+1
    scat = plt.scatter(xy_data[0,:],xy_data[1,:],marker='o',c='r',s=5,zorder=10)

x = np.linspace(-10,10,500)
fig = plt.figure
fig, axs = plt.subplots(1, 2, sharey=True, tight_layout=True)
axs[0].hist(v, bins=100)
plt.plot(x,np.exp(-x**2*betha/2)*9200)


Z = calZalternative(pointnum,V_backup,x_min,x_max,y_min,y_max)
Z_res = calcZinResPhaseSpace(pointnum,V_backup,x_min,x_max,y_min,y_max)
P = calcP(deltaE,betha,Z,mainsaddlepos)
predtrans = P*traj_num_max*(times - buffer)