# -*- coding: utf-8 -*-
"""
Created on Tue Jun  2 10:23:39 2015

@author: nicolas
"""

import numpy as np
import matplotlib.pyplot as plt




def Potential(x,y):
    V0 = 1.6022*10**(-19) # 1eV in kg*m^(2)s^(-2)
    a = 3. * 10**(-10)    # Gitterkonstante in m (=3A)
    Energy = 0.25 * V0 * ( 2. + np.cos(2*np.pi/a*x) + np.cos(2*np.pi/a*y)) #Potenzial V, siehe Aufgabenblatt
    Force = np.array([0.25 * V0 * 2.*np.pi / a * np.sin(2.*np.pi/a*x), 0.25 * V0 * 2.*np.pi / a * np.sin(2.*np.pi / a * y)]) #Gradient: -dEnergy/dx, -dEnergy/dy
    return [Energy, Force]




def Verlet(r1, r0, dt):             # Verlet(Positionsa[:,i-1],Positionsa[:,i-2], dt)
    m = 26.98*1.66*10**(-27)        # Masse Al-Atom in kg
    F = Potential(r1[0], r1[1])[1]  # Kraft bei Step (i-1)
    
    nextpos = 2.*r1 - r0 + dt**2 * F/m
    Ekin= 0.5*m*(np.linalg.norm((nextpos-r0)/(2.*dt))**2)  # Ekin am der der aktuellen Position (i-1), also nicht an nextpos (i)
    
    return nextpos, Ekin


def Initialize(Ekin, angle, Stepnumber, dt):
    posstart = np.array([1.5, 1.5])*10**(-10)
    Ekin = Ekin*1.6022*10**(-19) # Umrechnung in Joule
    m = 26.98*1.66*10**(-27)     # Masse Al-Atom in kg
    ForceStart =  Potential(posstart[0], posstart[1])[1] # in erstes Listenelement der ForceStart-Liste wird Energie und Kraft eingetragen
    v = np.sqrt(2*Ekin/m)        # Startgeschwindigkeit
    
    vstart = np.array([v*np.cos(float(angle)/180*np.pi), v*np.sin(float(angle)/180*np.pi)]) # Startgeschwindigkeit als Vektor
    
    pos1 = posstart + dt*vstart + 0.5*ForceStart/m*dt**2 # Simple Integration der Newtonschen Bewegungsgleichung
    
    posArray=np.zeros([2, Stepnumber]) # Nullmatrix der Größe '2 x Stepnumber' erstellen
    Etota = np.zeros(Stepnumber)       # Nullvektor der Größe 'Stepnumber' erstellen
    Ekina = np.zeros(Stepnumber)
    posArray[:,0] = posstart
    posArray[:,1] = pos1
    
    E1 = Potential(pos1[0], pos1[1])[0] # Epot nach dem ersten Schritt
    
    Etota[0] = Potential(posstart[0], posstart[1])[0] + Ekin # Epot+Ekin
    Etota[1] = Potential(posstart[0], posstart[1])[0] + Ekin # Energieerhaltung
    Ekina[0] = Ekin
    Ekina[1] = Etota[1] - E1
    return posArray, Etota, Ekina



##################################
#The real programme starts here


Stepnumber = 1000 
dt = 10.**(-15)   # in s
angle = 35        # in degrees
Ekin = 0.85       # in eV


[Positionsa, Etota, Ekina] = Initialize(Ekin,angle, Stepnumber, dt) # Nullvektoren werden erstellt mit Länge=Stepnumber. Erste und zweite Dimension werden berechnet

for i in range(Stepnumber):
    if i>1: #Pos 0 und 1 überspringen, die wurden bereits von Initialize berechnet
        Posout, Ekinout = Verlet(Positionsa[:,i-1],Positionsa[:,i-2], dt)  # Posout: Position an diesem step (i) | Ekinout: Ekin am letzten Step (i-1)
        Enerout = Potential(Positionsa[:,i-1][0], Positionsa[:,i-1][1])[0] # Energout : Epot am letzten step (i-1)
        Positionsa[:,i] = Posout
        Ekina[i-1] = Ekinout
        Etota[i-1] = Enerout+Ekinout



################################################################
#now this is just for the underlying potential contourplot

minx = np.min(Positionsa[0,:])
maxx = np.max(Positionsa[0,:])
miny = np.min(Positionsa[1,:])
maxy = np.max(Positionsa[1,:])


x = np.linspace(minx-10**-10, maxx+10**-10, 300)
y = np.linspace(miny-10**-10, maxy+10**-10, 300)

xv, yv = np.meshgrid(x, y)
z = Potential(xv, yv)[0]
plt.contourf(xv*10**10,yv*10**10,z)

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


plt.plot(Positionsa[0,:]*10**10, Positionsa[1,:]*10**10, lw = 5, c = "k", alpha =0.7)
plt.xlabel('x in Angstrom')
plt.ylabel(r'y in Angstrom')
plt.title("MD: Atom on surface")
plt.show()