import os
import time
import numpy as np

from ase import Atoms
from ase.parallel import paropen
from ase.units import Hartree, Bohr
from ase.io import Trajectory
from ase.calculators.singlepoint import SinglePointCalculator
from gpaw import GPAW
from gpaw.mpi import world
from gpaw.tddft import TDDFT
from gpaw.tddft.ehrenfest import EhrenfestVelocityVerlet


name = 'h2_osc'

# Equilibrium distance in Ang cf. setups page for H dimer
d_bond = 0.754  # ~0.766 during oscillation
d_disp = 0.03

# Timestep and expected oscillatory period in attoseconds
timestep = 5.0
period = 7.58e3  # ~545.7 meV cf. CRC Handbook of Phys. & Chem. #09_08_91

ndiv = int(np.ceil(0.1e3 / timestep))  # update stats every 0.1 fs
niter = ndiv * int(np.ceil(2 * period / (ndiv * timestep)))

if __name__ == '__main__':
    if not os.path.isfile(name + '_gs.gpw'):
        atoms = Atoms('H2', positions=[(0, 0, 0), (0, 0, d_bond + d_disp)])
        atoms.set_pbc(False)
        atoms.center(vacuum=4.0)
        cell_c = np.sum(atoms.get_cell()**2, axis=1)**0.5
        N_c = 8 * np.round(cell_c / (0.18 * 8))
        calc = GPAW(mode='fd',
                    gpts=N_c,
                    nbands=1,
                    basis='dzp',
                    txt=name + '_gs.txt',
                    symmetry={'point_group': False})
        atoms.calc = calc
        atoms.get_potential_energy()
        calc.write(name + '_gs.gpw', mode='all')
        del atoms, calc
        time.sleep(10)

    while not os.path.isfile(name + '_gs.gpw'):
        print('Node %d waiting for file...' % world.rank)
        time.sleep(10)
    world.barrier()

    tdcalc = TDDFT(name + '_gs.gpw', txt=name + '_td.txt', propagator='EFSICN')
    ehrenfest = EhrenfestVelocityVerlet(tdcalc)
    traj = Trajectory(name + '_td.traj', 'w', tdcalc.get_atoms())

    t0 = time.time()
    f = paropen(name + '_td.log', 'w')
    for i in range(1, niter + 1):
        ehrenfest.propagate(timestep)

        if i % ndiv == 0:
            rate = 60 * ndiv / (time.time() - t0)
            ekin = tdcalc.atoms.get_kinetic_energy()
            epot = tdcalc.get_td_energy() * Hartree
            F_av = ehrenfest.F * Hartree / Bohr
            s = 'i=%06d (%6.2f min^-1), ekin=%13.9f, epot=%13.9f, etot=%13.9f'
            print(s % (i, rate, ekin, epot, ekin + epot), file=f)
            t0 = time.time()

            # Hack to prevent calls to GPAW::get_potential_energy when saving
            spa = tdcalc.get_atoms()
            spc = SinglePointCalculator(spa, energy=epot, forces=F_av)
            spa.calc = spc
            traj.write(spa)
    f.close()
    traj.close()
