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

# ## Tuning the range separation
#
# This example tunes the range-separation parameter `gamma` for a long-range corrected functional in ADF. For each gamma value, it evaluates the J function that combines HOMO-IP differences for neutral and anionic states, then selects the gamma that minimizes J.
#

# ### Initial Imports

import psutil
from typing import List, Sequence, Tuple

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from scm.plams import (
    AMSJob,
    Atom,
    JobRunner,
    Molecule,
    MultiJob,
    Results,
    Settings,
    config,
    view,
)


# ### Define Helper Job and Result Classes
#
# A `GammaJob` runs three single-point calculations for one gamma value (charge -1, 0, +1), and `GammaResults` computes the corresponding J value.


class GammaResults(Results):
    @staticmethod
    def get_difference(job: AMSJob, job_plus: AMSJob) -> float:
        """Return HOMO + (E_plus - E) for one charge state pair."""
        homo = float(job.results.readrkf("Properties", "HOMO", file="engine"))
        ionization = float(job_plus.results.get_energy() - job.results.get_energy())
        return ionization + homo

    def get_J(self) -> float:
        n_term = GammaResults.get_difference(self.job.children[1], self.job.children[2])
        a_term = GammaResults.get_difference(self.job.children[0], self.job.children[1])
        return float((n_term**2 + a_term**2) ** 0.5)


class GammaJob(MultiJob):
    _result_type = GammaResults

    def __init__(
        self,
        molecule: Molecule,
        gamma: float,
        charge: int,
        spins: Tuple[int, int, int],
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.molecule = molecule
        self.gamma = gamma
        self.charge = charge
        self.spins = spins

    def prerun(self) -> None:
        charges = [self.charge - 1, self.charge, self.charge + 1]
        for charge, spin in zip(charges, self.spins):
            name = f"{self.name}_charge_{charge}".replace("-", "minus")
            child = AMSJob(name=name, molecule=self.molecule, settings=self.settings)
            child.molecule.properties.charge = charge
            child.settings.input.adf.xc.rangesep = f"gamma={self.gamma:f}"

            if spin != 0:
                child.settings.input.adf.unrestricted = True
                child.settings.input.adf.SpinPolarization = spin

            self.children.append(child)


def gamma_scan(
    gammas: Sequence[float],
    settings: Settings,
    molecule: Molecule,
    name: str = "scan",
    charge: int = 0,
    spins: Tuple[int, int, int] = (1, 0, 1),
) -> List[Tuple[float, float]]:
    jobs = [
        GammaJob(
            molecule=molecule,
            settings=settings,
            gamma=float(gamma),
            charge=charge,
            spins=spins,
            name=f"{name}_gamma_{gamma}",
        )
        for gamma in gammas
    ]

    results = [job.run() for job in jobs]
    j_values = [result.get_J() for result in results]
    return list(zip([float(g) for g in gammas], j_values))


def plot_scan(scan_results: Sequence[Tuple[float, float]]) -> Axes:
    fig, ax = plt.subplots(figsize=(6, 4))
    gammas = [gamma for gamma, _ in scan_results]
    j_values = [j for _, j in scan_results]
    ax.plot(gammas, j_values, marker="o")
    ax.set_xlabel("gamma")
    ax.set_ylabel("J")
    fig.tight_layout()
    return ax


# ### Run Gamma Scan
#
# Configure PLAMS to run independent gamma points in parallel.

config.default_jobrunner = JobRunner(parallel=True, maxjobs=psutil.cpu_count(logical=False))


# Set up ADF single-point settings with PBE and XCfun enabled.
#

s = Settings()
s.input.ams.task = "SinglePoint"
s.input.adf.basis.type = "DZP"
s.input.adf.basis.core = "None"
s.input.adf.xc.gga = "PBE"
s.input.adf.xc.xcfun = True
s.runscript.nproc = 1


# Build a toy H2 molecule for the scan.
#

mol = Molecule()
mol.add_atom(Atom(symbol="H", coords=(0.0, 0.0, -0.3540)))
mol.add_atom(Atom(symbol="H", coords=(0.0, 0.0, 0.3540)))
view(mol, width=250, height=250, direction="along_x", picture_path="picture1.png")


# Run a small gamma scan. For production use, use a wider range and finer spacing.
#

gammas = np.around(np.arange(1.2, 1.9, 0.1), decimals=3)
scan_results = gamma_scan(gammas, s, mol, name="gamma_scan")


###


print("== Results ==")
print("gamma	J")
for gamma, j_value in scan_results:
    print(f"{gamma:.4f}	{j_value:.8f}")

best_gamma = min(scan_results, key=lambda pair: pair[1])[0]
print(f"Optimal gamma value: {best_gamma:.4f}")


# Plot J as a function of gamma and identify the minimum.
#

ax = plot_scan(scan_results)
ax
ax.figure.savefig("picture2.png")