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)