"""Compute the ASE energy-volume curve from the official NequIP example."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

import numpy as np
import torch
from ase.build import bulk

from onescience.utils.nequip.integrations.ase import NequIPCalculator


def default_compiled_model() -> str | None:
    models_dir = os.environ.get("ONESCIENCE_MODELS_DIR")
    if not models_dir:
        return None
    return str(Path(models_dir) / "NequIP" / "NequIP-OAM-L-0.1.nequip.pth")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--compiled-model", default=default_compiled_model())
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--element", default="Si")
    parser.add_argument("--crystal-structure", default="diamond")
    parser.add_argument("--lattice-constant", type=float, default=5.43)
    parser.add_argument("--supercell", type=int, default=3)
    parser.add_argument("--scale-min", type=float, default=0.95)
    parser.add_argument("--scale-max", type=float, default=1.05)
    parser.add_argument("--num-points", type=int, default=10)
    parser.add_argument("--output", default="outputs/energy_volume.json")
    parser.add_argument("--plot", default="outputs/energy_volume.png")
    args = parser.parse_args()

    if not args.compiled_model:
        parser.error("--compiled-model is required when ONESCIENCE_MODELS_DIR is unset")
    compiled_model = Path(args.compiled_model).expanduser().resolve()
    if not compiled_model.is_file():
        parser.error(f"compiled model not found: {compiled_model}")
    if args.num_points < 2:
        parser.error("--num-points must be at least 2")
    if args.supercell < 1:
        parser.error("--supercell must be positive")

    calculator = NequIPCalculator.from_compiled_model(
        compile_path=str(compiled_model),
        chemical_species_to_atom_type_map={args.element: args.element},
        device=args.device,
    )

    points = []
    for scale in np.linspace(args.scale_min, args.scale_max, args.num_points):
        atoms = bulk(
            args.element,
            crystalstructure=args.crystal_structure,
            a=args.lattice_constant * float(scale),
            cubic=True,
        )
        atoms *= (args.supercell,) * 3
        atoms.calc = calculator
        energy = float(atoms.get_potential_energy())
        forces = atoms.get_forces()
        points.append(
            {
                "scale": float(scale),
                "volume_angstrom3": float(atoms.get_volume()),
                "energy_ev": energy,
                "energy_ev_per_atom": energy / len(atoms),
                "max_force_ev_per_angstrom": float(
                    np.linalg.norm(forces, axis=1).max()
                ),
            }
        )

    energies = np.asarray([point["energy_ev"] for point in points])
    volumes = np.asarray([point["volume_angstrom3"] for point in points])
    minimum_index = int(np.argmin(energies))
    result = {
        "compiled_model": str(compiled_model),
        "device": args.device,
        "device_name": torch.cuda.get_device_name(0)
        if args.device.startswith("cuda") and torch.cuda.is_available()
        else "cpu",
        "element": args.element,
        "crystal_structure": args.crystal_structure,
        "base_lattice_constant_angstrom": args.lattice_constant,
        "supercell": [args.supercell] * 3,
        "num_atoms": len(atoms),
        "points": points,
        "sampled_minimum": points[minimum_index],
    }

    output_path = Path(args.output).expanduser().resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")

    if args.plot:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        plot_path = Path(args.plot).expanduser().resolve()
        plot_path.parent.mkdir(parents=True, exist_ok=True)
        plt.figure(figsize=(8, 6))
        plt.plot(volumes, energies, marker="o", label="E-V Curve")
        plt.xlabel("Volume (Angstrom^3)", fontsize=14)
        plt.ylabel("Energy (eV)", fontsize=14)
        plt.title(f"Energy-Volume Curve for Cubic {args.element}", fontsize=16)
        plt.legend(fontsize=12)
        plt.grid()
        plt.tight_layout()
        plt.savefig(plot_path, dpi=160)
        plt.close()
        result["plot"] = str(plot_path)
        output_path.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")

    print(f"points: {len(points)}")
    print(f"atoms per point: {result['num_atoms']}")
    print(f"sampled minimum: {result['sampled_minimum']}")
    print(f"result: {output_path}")


if __name__ == "__main__":
    main()