# Workflow with PLAMS

# A couple of useful function for extracting results:

@add_to_class(AMSResults)
def get_excitations(results):
    """Returns excitation energies (in eV) and oscillator strenghts (in Debye)."""
    if results.job.ok():
        exci_energies_au = results.readrkf('Excitations SS A','excenergies',file='engine')
        oscillator_str_au = results.readrkf('Excitations SS A','oscillator strengths',file='engine')
        # The results are stored in atomic units. Convert them to more convenient units:
        exci_energies = Units.convert(exci_energies_au,'au','eV')
        oscillator_str = Units.convert(oscillator_str_au,'au','Debye')
        return exci_energies, oscillator_str
    else:
        return [],[]


@add_to_class(AMSResults)
def has_good_excitations(results, min_energy, max_energy, oscillator_str_threshold=1E-4):
    """Returns True if there is at least one excitation with non-vanishing oscillator strenght
       in the energy range [min_energy, max_energy]. Unit for min_energy and max energy: eV. """
    exci_energies, oscillator_str = results.get_excitations()
    for e,o in zip(exci_energies, oscillator_str):
        if min_energy < e < max_energy and o > oscillator_str_threshold:
            return True
    return False


# Calculation settings:
# =====================

# Settings for geometry optimization with the AMS driver:
go_sett = Settings()
go_sett.input.ams.Task = 'GeometryOptimization'
go_sett.input.ams.GeometryOptimization.Convergence.Gradients = 1.0e-4

# Settings for single point calculation with the AMS driver
sp_sett = Settings()
sp_sett.input.ams.Task = 'SinglePoint'

# Settings for the DFTB engine (including excitations)
dftb_sett = Settings()
dftb_sett.input.dftb.Model = 'SCC-DFTB'
dftb_sett.input.dftb.ResourcesDir = 'QUASINANO2015'
dftb_sett.input.dftb.Properties.Excitations.TDDFTB.calc = 'singlet'
dftb_sett.input.dftb.Properties.Excitations.TDDFTB.lowest = 10
dftb_sett.input.dftb.Occupation.Temperature = 5.0

# Settings for the geometry optimization with the ADF engine
adf_sett = Settings()
adf_sett.input.adf.Basis.Type = 'DZP'
adf_sett.input.adf.NumericalQuality = 'Basic'

# Settings for the excitation calculation using the ADF engine
adf_exci_sett = Settings()
adf_exci_sett.input.adf.Basis.Type = 'TZP'
adf_exci_sett.input.adf.XC.GGA = 'PBE'
adf_exci_sett.input.adf.NumericalQuality = 'Basic'
adf_exci_sett.input.adf.Symmetry = 'NoSym'
adf_exci_sett.input.adf.Excitations.lowest = 10
adf_exci_sett.input.adf.Excitations.OnlySing = ''


# Import all xyz files in the folder 'molecules'

molecules = read_molecules('molecules')


print("Step 1: prescreening with DFTB")
print("==============================")

promising_molecules = {}

for name, mol in molecules.items():
  dftb_job = AMSJob(name='DFTB_'+name, molecule=mol, settings=go_sett+dftb_sett)
  dftb_job.run()

  if dftb_job.results.has_good_excitations(1,6):
    promising_molecules[name] = dftb_job.results.get_main_molecule()

print(f"Found {len(promising_molecules)} promising molecules with DFTB")


print("Step 2: Optimization and excitations calculation with ADF")
print("=========================================================")

for name, mol in promising_molecules.items():
  adf_go_job = AMSJob(name='ADF_GO_'+name, molecule=mol, settings=go_sett+adf_sett)
  adf_go_job.run()

  optimized_mol = adf_go_job.results.get_main_molecule()

  adf_exci_job = AMSJob(name='ADF_exci_'+name, molecule=optimized_mol, 
                        settings=sp_sett+adf_exci_sett)
  adf_exci_job.run()

  if adf_exci_job.results.has_good_excitations(2,4):
    print(f"Molecule {name} has excitation(s) satysfying our criteria!")
    print(optimized_mol)
    exci_energies, oscillator_str = adf_exci_job.results.get_excitations()
    print("Excitation energy [eV], oscillator strength:")
    for e,o in zip(exci_energies, oscillator_str):
        print(f"{e:8.4f}, {o:8.4f}")