# %%
"""
## Extra exercise - vibrational energy
The energy calculated with DFT is the electronic energy at 0K. However, to model catalytic reactions we are usually intereseted in the energy landscape at finite temperature. In this exercise we will calculate the energy contributions from vibrations and see how they affect the splitting of N<sub>2</sub> on Ru.

We calculate the vibrational energy in the harmonic approximation using the finite displacement method. For further reading see for example:

* [Stoffel et al, Angewandte Chemie Int. Ed., 49, 5242-5266 (2010)](https://onlinelibrary.wiley.com/doi/abs/10.1002/anie.200906780)
* [Togo and Tanaka, Scripta Materialia 108, 1-5 (2015)](https://www.sciencedirect.com/science/article/pii/S1359646215003127)

### Vibrational energy of the initial and final states
a) Calculating vibrations requires tighter convergence than normal energy calculations. Therefore you should first take your already optimised initial and final state geometries from the NEB calculations and relax them further to `fmax=0.01eV/Å` with the QuasiNewton optimiser and an energy cutoff of 450eV. Converge the eigenstates to 1e-8. (Note that for other systems you might need even tighter convergence!)

Submit the structures to the queue. The optimisation should take 10-15 mins for each structure on 8 cores.
"""

# %%
# teacher:
# Script for reading previous structure and converging
from gpaw import GPAW, PW
from ase.constraints import FixAtoms
from ase.optimize import QuasiNewton
import numpy as np
from ase.io import read, write

for name in ['N2Ru.traj', '2Nads.traj']:
    slab = read(name)

    z = slab.positions[:, 2]
    constraint = FixAtoms(mask=(z < z.min() + 1.0))
    slab.set_constraint(constraint)

    calc = GPAW(xc='PBE',
                mode=PW(450),
                txt='tight-' + name[:-4] + 'txt',
                kpts={'size': (4, 4, 1), 'gamma': True},
                convergence={'eigenstates': 1e-8})
    slab.calc = calc

    dyn = QuasiNewton(slab,
                      trajectory='tight-' + name,
                      maxstep = 0.02,
                      logfile='tight-' + name[:-4] + 'log')
    dyn.run(fmax=0.01)

# %%
"""
b) Once you have done this you can calculate the vibrations using the [vibrations module in ASE](https://wiki.fysik.dtu.dk/ase/ase/vibrations/vibrations.html) following the template below. We only calculate the vibrations of the adsorbate and assume that the frequencies of the substrate are unchanged - this is a common assumption. Use 4 displacements to fit the frequencies and the same calculator parameters as in a).
Submit the calculations for the initial and the final state to the queue. It will take a while to run, but you can start preparing your analysis (part c and d) or read some of the references in the meantime.
"""

# %%
# You can't run this cell - use it as a starting point for your python script!
from ase.io import read
from ase.vibrations import Vibrations
from gpaw import GPAW, PW

slab = read('tight-N2Ru.traj')  # student: slab = read('path_to_your_structure')
calc = GPAW(xc='PBE',
            mode=PW(450),  # student: mode=PW(450),
            kpts={'size': (4, 4, 1), 'gamma': True},  # student: kpts=?,
            convergence={'eigenstates': 1e-8},  # student; convergence=?,
            symmetry={'point_group': False},
            txt='vib.txt')
slab.calc = calc
Uini = slab.get_potential_energy()

vib = Vibrations(slab,
                 name='vib',
                 indices=[8, 9],  # student: indices=[?, ?],
                 nfree=4)
vib.run()
vib.summary(log='vib_summary.log')
for i in range(6):
    vib.write_mode(i)

# %%
from ase.io import read
from ase.vibrations import Vibrations
from gpaw import GPAW, PW

slab = read('tight-2Nads.traj')
calc = GPAW(xc='PBE',
            mode=PW(450),  # student: mode=PW(450),
            kpts={'size': (4, 4, 1), 'gamma': True},  # student: kpts=?,
            convergence={'eigenstates': 1e-8},  # student; convergence=?,
            symmetry={'point_group': False},
            txt='vib2.txt')
slab.calc = calc
Ufin = slab.get_potential_energy()

vib2 = Vibrations(slab,
                  name='vib2',
                  indices=[8, 9],  # student: indices=[?, ?],
                  nfree=4)
vib2.run()
vib2.summary(log='vib2_summary.log')
for i in range(6):
    vib2.write_mode(i)

# %%
"""
The final three lines write out the frequencies and `.traj` files with animations of the different modes in order of their energy. Take a look at the vibrational modes in the files. Do you understand why the first mode has a low energy, while the last one has a high energy?
"""

# %%
"""
c) Analyse the frequencies in the harmonic approximation:
"""

# %%
from ase.thermochemistry import HarmonicThermo

T = 300  # Kelvin
energies = [0.01, 0.05, 0.10]  # An example only - insert your calculated energy levels here - in eV!
vibs = HarmonicThermo(energies)
Efree = vibs.get_helmholtz_energy(T, verbose=True)
print('Free energy at 300 K: ', Efree)

# %%
"""
The `verbose` keyword gives a detailed description of the different contributions to the free energy. For more information on what the different contributions are see the [ASE background webpage](https://wiki.fysik.dtu.dk/ase/ase/thermochemistry/thermochemistry.html#background) (go to the **Harmonic limit** sub-heading).
Now try to calculate how the different contributions change with temperature. You can for example make a `for` loop and use the `get_entropy()` and `get_internal_energy()` methods [(see description here)](https://wiki.fysik.dtu.dk/ase/ase/thermochemistry/thermochemistry.html#ase.thermochemistry.IdealGasThermo.get_enthalpy).
"""

# %%
"""
d) Calculate how the vibrational energy affects the overall reaction energy. How does it change with temperature? Which contribution is important for the change in reaction energy?
"""

# %%
# teacher:
# Script for analysing energy contributions from vibrations
from ase.io import read, write
from ase.thermochemistry import HarmonicThermo
import numpy as np
from matplotlib import pyplot as plt

energies = vib.get_energies()  # initial vib. energy levels
energiesfin = vib2.get_energies()  # final vib. energy levels

# The actual analysis
ads = HarmonicThermo(energies)
adsfin = HarmonicThermo(energiesfin)
print('DFT reaction energy: ', Ufin - Uini)
Sini = []
Sfin = []
U = []
U2 = []
temp = [10, 50, 100, 150, 200, 250, 300, 350, 400, 450, 500]
for T in temp:
    Sini.append(-T * ads.get_entropy(T, verbose=False))
    U.append(ads.get_internal_energy(T, verbose=False))
    Sfin.append(-T * adsfin.get_entropy(T, verbose=False))
    U2.append(adsfin.get_internal_energy(T, verbose=False))
    Efree = ads.get_helmholtz_energy(T, verbose=False)
    Efreefin = adsfin.get_helmholtz_energy(T, verbose=False)
    print(f'Reaction free energy at {T} K: ',
          Ufin + Efreefin - (Uini + Efree))
plt.plot(temp, Sini, 'r', label='-T*S ini')
plt.plot(temp, U, 'b', label='U ini')
plt.plot(temp, Sfin, 'r--', label='-T*S fin')
plt.plot(temp, U2, 'b--', label='U fin')
plt.plot(temp, np.array(U) + np.array(Sini), 'm', label='Efree ini')
plt.plot(temp, np.array(U2) + np.array(Sfin), 'm--', label='Efree fin')
plt.legend()
plt.savefig('vib.png')

# %%
"""
e) To make sure that your NEB is converged you should also calculate the vibrational energy of the transition state. Again, this requires tighter convergence than we have used in the NEB exercise. This takes a while to run so to save time, we provide the transition state geometry from a reasonably converged NEB (i.e. `fmax=0.01`, a cutoff energy of 500eV and eigenstates converged to 1e-8) in the file `TS.xyz`. Calculate the vibrations with these parameters. How many imaginary modes do you get and how do they look? What does this mean?
"""

# %%
# teacher:
from ase.optimize import BFGS
from ase.neb import NEB
from ase.io import Trajectory

initial = read('N2Ru.traj')
final = read('2Nads.traj')

images = read('neb.traj@-5:-2')
constraint = FixAtoms(indices=list(range(4)))
ts = images[1]
ts.set_constraint(constraint)
calc = GPAW(xc='PBE',
            mode=PW(500),
            txt='ts.txt',
            kpts={'size': (4, 4, 1), 'gamma': True},
            convergence={'eigenstates': 1e-8})
ts.calc = calc

neb = NEB(images, k=1.0, climb=True)
qn = BFGS(neb, logfile='neb.log')
traj = Trajectory('ts.traj', 'w', ts)
qn.attach(traj)
qn.run(fmax=0.01)
write('TS.xyz', ts)


# %%
# teacher:
ts.calc = GPAW(xc='PBE',
               mode=PW(500),
               txt='vibts.txt',
               kpts={'size': (4, 4, 1), 'gamma': True},
               symmetry={'point_group': False},
               convergence={'eigenstates': 1e-8})
vib = Vibrations(ts, name='vibts', indices=(8, 9), nfree=4, delta=0.02)
vib.run()
vib.summary(log='vibts_summary.log')
for i in range(6):
    vib.write_mode(i)

#---------------------
#  #    meV     cm^-1
#---------------------
#  0   73.7i    594.2i
#  1    7.8      63.0
#  2   39.6     319.6
#  3   46.0     371.1
#  4   67.9     547.4
#  5   68.8     554.8
#---------------------
#Zero-point energy: 0.115 eV
#The imaginary mode is beautifully along the reaction coordinate!
