#!/usr/bin/env python
# coding: utf-8

# ## Reduction and oxidation potentials
#
# This example runs the three PLAMS redox recipes for benzoquinone in water and compares their computed oxidation/reduction potentials.
#

# ### Import PLAMS recipe classes and utilities.
#

from typing import Iterable

from scm.plams import CRSJob, Settings, from_smiles, view
from scm.plams.recipes.redox import (
    AMSRedoxDirectJob,
    AMSRedoxScreeningJob,
    AMSRedoxThermodynamicCycleJob,
)


# ### Build benzoquinone and define solvent settings.
#

molecule = from_smiles("C1=CC(=O)C=CC1=O", forcefield="uff")
view(molecule, width=260, height=220)
solvent_name = "Water"
solvent_coskf = CRSJob.database() + "/Water.coskf"
reduction = True
oxidation = True
vibrations = True


# ### Define ADF settings used by the direct and thermodynamic-cycle workflows.
#

adf_settings = Settings()
adf_settings.input.adf.basis.type = "DZP"
adf_settings.input.adf.xc.gga = "PBE"
adf_settings.input.ams.GeometryOptimization.Convergence.Quality = "Basic"


# ### Create helper printing function and instantiate jobs.


def print_results(jobs: Iterable) -> None:
    she = 4.42  # eV, absolute SHE scale
    print("The experimental reduction potential of benzoquinone is +0.10 V vs. SHE")
    print(
        f"{'Jobname':24s} {'Eox(vib,rel-to-SHE)[V]':24s} {'Ered(vib,rel-to-SHE)[V]':24s} "
        f"{'Eox(rel-to-SHE)[V]':24s} {'Ered(rel-to-SHE)[V]':24s}"
    )

    for job in jobs:
        row = f"{job.name:24s}"
        for vib in [True, False]:
            try:
                eox = f"{job.results.get_oxidation_potential(vibrations=vib) - she:.2f}"
            except Exception:
                eox = "N/A"
            row += f" {eox:24s}"

            try:
                ered = f"{job.results.get_reduction_potential(vibrations=vib) - she:.2f}"
            except Exception:
                ered = "N/A"
            row += f" {ered:24s}"
        print(row)


jobs = [
    AMSRedoxScreeningJob(
        name="quick_screening",
        molecule=molecule,
        reduction=reduction,
        oxidation=oxidation,
        solvent_coskf=solvent_coskf,
    ),
    AMSRedoxDirectJob(
        name="direct_best_method",
        molecule=molecule,
        settings=adf_settings,
        reduction=reduction,
        oxidation=oxidation,
        vibrations=vibrations,
        solvent=solvent_name,
    ),
    AMSRedoxThermodynamicCycleJob(
        name="thermodynamic_cycle",
        molecule=molecule,
        settings=adf_settings,
        reduction=reduction,
        oxidation=oxidation,
        vibrations=vibrations,
        solvent=solvent_name,
    ),
]


# ### Run all workflows and print per-workflow plus summary tables

for job in jobs:
    job.run()

for job in jobs:
    job.ok()


print("Final summary:")
print_results(jobs)