import sys
import numpy as np

config.log.stdout = 0

# This test calculates the vibrational fine structure of the S_0 -> S_1 transition in anthracene.

molfile = "anthracene.xyz"
excit = 1

# Common settings for all DFTB calculations:
comin = Settings()
comin.input.DFTB.resourcesdir = "DFTB.org/3ob-freq-1-2"
comin.input.DFTB.model = "DFTB3"


# ========= auxilliary functions ==================================================================
def get_total_energy(results):
    nprop = results.readrkf("Properties", "nEntries", file="dftb")
    for i in range(1, nprop + 1):
        if results.readrkf("Properties", "Subtype(%i)" % i, file="dftb").strip() == "DFTB Final Energy":
            return results.readrkf("Properties", "Value(%i)" % i, file="dftb")
    return None


def get_zero_point_energy(results):
    freqs = results.readrkf("Vibrations", "Frequencies[cm-1]", file="dftb")
    if isinstance(freqs, list):
        return Units.convert(0.5 * sum(freqs), "cm^-1", "Hartree")
    else:
        return Units.convert(0.5 * freqs, "cm^-1", "Hartree")


def extract_spectrum(fcf_results):
    return np.array(fcf_results.readkf("Fcf", "spectrum")).reshape(2, -1).transpose()


# ========= STEP 1: Ground state ==================================================================

# Optimize ground state geometry:
gs_mol_unoptimized = Molecule(filename=molfile)
gs_go = AMSJob(name="gs_go", molecule=gs_mol_unoptimized, settings=comin)
gs_go.settings.input.ams.Task = "GeometryOptimization"
gs_go.settings.input.ams.GeometryOptimization.convergence = "Gradients=1.0e-5"
gs_go_results = gs_go.run()
if not gs_go.check():
    print("ERROR: Ground state optimization crashed")
    sys.exit(1)
if gs_go_results.grep_output("Optimization Did Not Converge"):
    print("ERROR: Ground state optimization did not converge")
    sys.exit(1)
gs_mol_optimized = gs_go_results.get_molecule("Molecule")

# Calculate frequencies and normal modes of the ground state:
gs_freq = AMSJob(name="gs_freq", molecule=gs_mol_optimized, settings=comin)
gs_freq.settings.input.ams.Task = "SinglePoint"
gs_freq.settings.input.ams.properties.NormalModes = "true"
gs_freq.settings.input.ams.NumericalDifferentiation.Parallel.nCoresPerGroup = 1
gs_freq_results = gs_freq.run()
if not gs_freq.check():
    print("ERROR: Ground state frequency calculation crashed")
    sys.exit(1)

# Calculate vertical excitations:
gs_excit = AMSJob(name="gs_excit", molecule=gs_mol_optimized, settings=comin)
gs_excit.settings.input.ams.Task = "SinglePoint"
gs_excit.settings.input.DFTB.properties.excitations.tddftb.calc = "singlet"
gs_excit.settings.input.DFTB.properties.excitations.tddftb.lowest = excit + 9
gs_excit.settings.input.DFTB.properties.excitations.tddftb["print"] = "evcontribs"
gs_excit_results = gs_excit.run()
if not gs_excit.check():
    print("ERROR: Ground state excitations calculation crashed")
    sys.exit(1)

# Print ground state energies:
print("Energies in the ground state equilibrium geometry:")
E_DFTB_RGS = get_total_energy(gs_excit_results)
E_ZPE_RGS = get_zero_point_energy(gs_freq_results)
Delta_RGS = gs_excit_results.readrkf("Excitations SS A", "excenergies", file="dftb")[excit - 1]
E_GS = E_DFTB_RGS + E_ZPE_RGS
print("  E_DFTB(R_GS) = %f eV" % (Units.convert(E_DFTB_RGS, "Hartree", "eV")))
print("   E_ZPE(R_GS) = %f eV" % (Units.convert(E_ZPE_RGS, "Hartree", "eV")))
print("         E_GS  = %f eV" % (Units.convert(E_GS, "Hartree", "eV")))


# ========= STEP 2: Excited state =================================================================

# Optimize the excited state geometry:
ex_go = AMSJob(name="ex_go", molecule=gs_mol_optimized, settings=comin)
ex_go.settings.input.ams.Task = "GeometryOptimization"
ex_go.settings.input.ams.GeometryOptimization.convergence = "Gradients=1.0e-5"
ex_go.settings.input.DFTB.properties.excitations.tddftb.calc = "singlet"
ex_go.settings.input.DFTB.properties.excitations.tddftb.lowest = excit
ex_go.settings.input.DFTB.properties.excitations.tddftb["print"] = "evcontribs"
ex_go.settings.input.DFTB.properties.excitations.tddftbgradients.excitation = excit
ex_go.settings.input.DFTB.properties.excitations.tddftbgradients.eigenfollow = "true"
ex_go.settings.input.ams.log.info = "TDDFTBExcitationFollowerModule"
ex_go_results = ex_go.run()
if not ex_go.check():
    print("ERROR: Excited state optimization crashed")
    sys.exit(1)
if ex_go_results.grep_output("Optimization Did Not Converge"):
    print("ERROR: Excited state optimization did not converge")
    sys.exit(1)
ex_mol_optimized = ex_go_results.get_molecule("Molecule")

# Check if the potential energy surface was switched during the optimization:
# (This happens if the optimizer goes through a conical intersection.)
PES_switches = ex_go_results.grep_file("ams.log", "TD-DFTB Eigenfollower switching PES:")
if PES_switches:
    newexcit = int(PES_switches[-1].split()[-1])
    print("PES switched during EXGO!!! %i -> %i" % (excit, newexcit))
else:
    newexcit = excit

# Calculate frequencies and normal modes of the excited state:
ex_freq = AMSJob(name="ex_freq", molecule=ex_mol_optimized, settings=comin)
ex_freq.settings.input.ams.Task = "SinglePoint"
ex_freq.settings.input.ams.properties.NormalModes = "true"
ex_freq.settings.input.ams.NumericalDifferentiation.Parallel.nCoresPerGroup = 1
ex_freq.settings.input.DFTB.properties.excitations.tddftb.calc = "singlet"
ex_freq.settings.input.DFTB.properties.excitations.tddftb.lowest = newexcit
ex_freq.settings.input.DFTB.properties.excitations.tddftb["print"] = "evcontribs"
ex_freq.settings.input.DFTB.properties.excitations.tddftbgradients.excitation = newexcit
ex_freq_results = ex_freq.run()
if not ex_freq.check():
    print("ERROR: Excited state frequency calculation crashed")
    sys.exit(1)

# Calculate vertical excitations in excited state geometry:
ex_excit = AMSJob(name="ex_excit", molecule=ex_mol_optimized, settings=comin)
ex_excit.settings.input.ams.Task = "SinglePoint"
ex_excit.settings.input.DFTB.properties.excitations.tddftb.calc = "singlet"
ex_excit.settings.input.DFTB.properties.excitations.tddftb.lowest = newexcit + 9
ex_excit.settings.input.DFTB.properties.excitations.tddftb["print"] = "evcontribs"
ex_excit_results = ex_excit.run()
if not ex_excit.check():
    print("ERROR: Excited state geometry excitations calculation crashed")
    sys.exit(1)

# Print excited state energies:
print("Energies in the excited state equilibrium geometry:")
E_DFTB_REX = get_total_energy(ex_excit_results)
E_ZPE_REX = get_zero_point_energy(ex_freq_results)
Delta_REX = ex_excit_results.readrkf("Excitations SS A", "excenergies", file="dftb")[excit - 1]
E_EX = E_DFTB_REX + E_ZPE_REX + Delta_REX
print("  E_DFTB(R_EX) = %f eV" % (Units.convert(E_DFTB_REX, "Hartree", "eV")))
print("   E_ZPE(R_EX) = %f eV" % (Units.convert(E_ZPE_REX, "Hartree", "eV")))
print("   Delta(R_EX) = %f eV" % (Units.convert(Delta_REX, "Hartree", "eV")))
print("         E_EX  = %f eV" % (Units.convert(E_EX, "Hartree", "eV")))

# Print excitation energies:
print("Excitation energies:")
print("   Delta(R_GS) = %f eV" % (Units.convert(Delta_RGS, "Hartree", "eV")))
print("         E_0-0 = %f eV" % (Units.convert(E_EX - E_GS, "Hartree", "eV")))
print("          Diff = %f eV" % (Units.convert(Delta_RGS - (E_EX - E_GS), "Hartree", "eV")))


# ========= STEP 3: Vibrational fine structure with the FCF program ===============================

# Settings for the FCF program
fcfin = Settings()
fcfin.input.spectrum.spcmin = "0.0"
fcfin.input.spectrum.spcmax = "5000.0"
fcfin.input.spectrum.spclen = "501"
fcfin.input.spectrum.lineshape = "Stick"
fcfin.input.numericalquality = "Basic"
fcfin.input.translate = True
fcfin.input.rotate = True

# Calculate vibrational fine structure
fcf = FCFJob(
    name="fcf",
    settings=fcfin,
    inputjob1=gs_freq_results.rkfpath(file="dftb"),
    inputjob2=ex_freq_results.rkfpath(file="dftb"),
)
fcf_results = fcf.run()
if not fcf.check():
    print("ERROR: FCF calculation failed")
    sys.exit(1)

# Extract and print the spectrum:
spectrum = extract_spectrum(fcf_results)
np.set_printoptions(formatter={"float": " {: 0.8f} ".format}, threshold=1e6)
print("Vibrational fine structure:")
print("Energy [cm^-1]    Intensity")
print(spectrum)