# import libraries and set up the molecule geometry

from ase.units import Ry, eV, Ha
from ase.calculators.siesta import Siesta
from ase import Atoms
import numpy as np
import matplotlib.pyplot as plt

from ase.build import molecule

CH4 = molecule("CH4")

# visualization of the particle
from ase.visualize import view
view(CH4, viewer='x3d')

"""
Explanation: Easy Ab initio calculation with ASE-Siesta-Pyscf
No installation necessary, just download a ready to go container for any system, or run it into the cloud
We first import the necessary libraries and define the system using ASE
End of explanation
"""


# enter siesta input and run siesta
siesta = Siesta(
    mesh_cutoff=150 * Ry,
    basis_set='DZP',
    pseudo_qualifier='lda',
    energy_shift=(10 * 10**-3) * eV,
    fdf_arguments={
        'SCFMustConverge': False,
        'COOP.Write': True,
        'WriteDenchar': True,
        'PAO.BasisType': 'split',
        'DM.Tolerance': 1e-4,
        'DM.MixingWeight': 0.1,
        'MaxSCFIterations': 300,
        'DM.NumberPulay': 4,
        'XML.Write': True})

CH4.set_calculator(siesta)
e = CH4.get_potential_energy()

"""
Explanation: We can then run the DFT calculation using Siesta
End of explanation
"""


# compute polarizability using pyscf-nao

siesta.pyscf_tddft(label="siesta", jcutoff=7, iter_broadening=0.15/Ha,
        xc_code='LDA,PZ', tol_loc=1e-6, tol_biloc=1e-7, freq = np.arange(0.0, 15.0, 0.05))

# plot polarizability with matplotlib
%matplotlib inline

fig = plt.figure(1)
ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122)
ax1.plot(siesta.results["freq range"], siesta.results["polarizability nonin"][:, 0, 0].imag)
ax2.plot(siesta.results["freq range"], siesta.results["polarizability inter"][:, 0, 0].imag)

ax1.set_xlabel(r"$\omega$ (eV)")
ax2.set_xlabel(r"$\omega$ (eV)")

ax1.set_ylabel(r"Im($P_{xx}$) (au)")
ax2.set_ylabel(r"Im($P_{xx}$) (au)")

ax1.set_title(r"Non interacting")
ax2.set_title(r"Interacting")

fig.tight_layout()

"""
Explanation: The TDDFT calculations with PySCF-NAO
End of explanation
"""


