#!/usr/bin/env amspython
from __future__ import annotations

from scm.plams import AMSJob, AMSNVTJob, finish, init, log

from common import (
    SAMPLING_FREQUENCY,
    TARGET_TEMPERATURE_K,
    TIMESTEP_FS,
    completed_jobs,
    equilibration_metrics,
    log_equilibration_check,
    restart_settings,
    reaxff_settings,
    sorted_cumulative_jobs,
)


INCREMENT_PS = 5
MAXIMUM_PS = 100
STEPS_PER_INCREMENT = int(INCREMENT_PS * 1000 / TIMESTEP_FS)


def run_first_segment(setup_job: AMSJob) -> AMSJob:
    job = AMSNVTJob(
        name="eq_005ps",
        molecule=setup_job.results.get_main_molecule(),
        settings=reaxff_settings(),
        nsteps=STEPS_PER_INCREMENT,
        timestep=TIMESTEP_FS,
        samplingfreq=SAMPLING_FREQUENCY,
        temperature=TARGET_TEMPERATURE_K,
        thermostat="Berendsen",
        tau=100.0,
        calcpressure=True,
        writevelocities=True,
        writebonds=True,
        writemolecules=True,
    )
    log("Starting the first 5 ps Berendsen NVT equilibration segment from optimized coordinates.")
    job.run()
    if not job.ok():
        raise RuntimeError(f"Equilibration segment failed: {job.results.get_errormsg()}")
    log(f"Completed the cumulative 5 ps equilibration trajectory at {job.path}.")
    return job


def extend(previous_job: AMSJob, target_ps: int) -> AMSJob:
    target_steps = int(target_ps * 1000 / TIMESTEP_FS)
    settings = restart_settings(previous_job, target_steps)
    job = AMSJob(name=f"eq_{target_ps:03d}ps", molecule=None, settings=settings)
    log(f"Extending equilibration exactly from {previous_job.name} to {target_ps} ps with Restart and "
        "CopyRestartTrajectory Yes.")
    job.run()
    if not job.ok():
        raise RuntimeError(f"Equilibration extension to {target_ps} ps failed: {job.results.get_errormsg()}")
    log(f"Completed the cumulative {target_ps} ps equilibration trajectory at {job.path}.")
    return job


def main() -> None:
    setup_jobs = completed_jobs("01-initial-system_workdir", "setup_opt")
    if not setup_jobs:
        raise RuntimeError("No completed setup_opt job. Run 01-initial-system.py first.")
    setup_job = setup_jobs["setup_opt"]

    known = completed_jobs("02-equilibration_workdir", "eq_*ps")
    cumulative = sorted_cumulative_jobs(known, "eq_")
    if cumulative:
        log(f"Found {len(cumulative)} reusable completed equilibration trajectories; latest is "
            f"{cumulative[-1][0]} ps at {cumulative[-1][1].path}.")
    else:
        first = run_first_segment(setup_job)
        cumulative = [(INCREMENT_PS, first)]

    streak = 0
    converged = False
    for duration_ps, job in cumulative:
        if duration_ps < 10:
            log(f"Skipping the equilibration convergence test at {duration_ps} ps; at least 10 ps is required.")
            continue
        metrics = equilibration_metrics(job)
        streak = streak + 1 if metrics["passed"] else 0
        log_equilibration_check(duration_ps, metrics, streak)
        if streak >= 2:
            converged = True
            log(f"Equilibration accepted at {duration_ps} ps after two consecutive passing checks.")
            break

    while not converged and cumulative[-1][0] < MAXIMUM_PS:
        target_ps = cumulative[-1][0] + INCREMENT_PS
        job = extend(cumulative[-1][1], target_ps)
        cumulative.append((target_ps, job))
        if target_ps < 10:
            continue
        metrics = equilibration_metrics(job)
        streak = streak + 1 if metrics["passed"] else 0
        log_equilibration_check(target_ps, metrics, streak)
        if streak >= 2:
            converged = True
            log(f"Equilibration accepted at {target_ps} ps after two consecutive passing checks.")

    if not converged:
        log("Equilibration reached the 100 ps cap without two consecutive passing checks. "
            "The final capped trajectory will be used for production and reported as not equilibrated by the test.")


if __name__ == "__main__":
    init(folder="02-equilibration_workdir")
    try:
        main()
    finally:
        finish()