# /// script
# requires-python = "==3.12.*"
# dependencies = [
#   "numpy==2.4.1",
#   "pymatgen==2025.10.7",
#   "tqdm==4.67.1",
# ]
# ///

# SPDX-FileCopyrightText: Copyright (C) 2026 Patrick J. Taylor
# SPDX-License-Identifier: GPL-3.0-or-later

import glob
import json
import shutil
from pathlib import Path

import numpy as np
from numpy.typing import NDArray
from pymatgen.io.vasp import Oszicar, Vasprun
from tqdm.auto import tqdm


def extract_fe_magmoms(path: str, n_fe: int) -> NDArray[np.float64]:
    """
    Extract the Fe local magnetic moments from a VASP OUTCAR.

    Parameters
    ----------
    path:
        Filepath to an OUTCAR.
    n_fe:
        The number of Fe atoms.

    Returns
    -------
    magmoms:
        A numpy array containing the extracted magnetic moments.
    """
    with open(path) as stream:
        lines = stream.readlines()

    magmom_table_indices = tuple(
        idx for idx in range(len(lines)) if "magnetization" in lines[idx].split()
    )
    start_idx = magmom_table_indices[-1] + 4

    # For all calculations that we apply this function to, the Fe atoms are first in
    # the POSCAR, hence the naive approach of just taking the first n_fe entries.
    magmom_table = lines[start_idx : start_idx + n_fe]
    magmoms = np.array([float(line.split()[-1]) for line in magmom_table])

    return magmoms


compositions = ("Li2FeSiO4", "LiFeSiO4", "FeSiO4")
temperatures = ("500K", "1000K")
frame_sets = ((0, 6881, 16602, 39366), (0, 2581, 8315, 12276))
polymorphs = ("inv_pmn21", "p21", "pbn21", "pmn21", "pmnb", "pn21")

print(
    "Extracting PDOS, magmoms, ICOBIs and Wannier centres/isosurfaces for x = 2, 1, 0"
)
for composition in tqdm(compositions):
    prefix = f"sequential_delithiation/{composition}"
    destination = f"extracted_data/sequential_delithiation/{composition}"
    Path(destination).mkdir(parents=True)

    magmoms = extract_fe_magmoms(f"{prefix}/lobster/OUTCAR", 2)
    np.save(f"{destination}/magmoms.npy", magmoms)

    shutil.copy(f"{prefix}/lobster/DOSCAR.lobster", destination)
    shutil.copy(f"{prefix}/lobster/ICOBILIST.lobster", destination)
    shutil.copy(f"{prefix}/lobster/POSCAR", destination)

    for wannier_centres in glob.iglob(
        f"{prefix}/wannier/oxidation_states/localisation/*/*.xyz"
    ):
        shutil.copy(wannier_centres, destination)

    for isosurface in glob.iglob(f"{prefix}/wannier/LUMO/localisation/*.xsf"):
        shutil.copy(isosurface, destination)

print(
    "Extracting energies, magmoms, ICOBIs and Wannier centres for AIMD trajectories (500 K, 1000 K)"
)
for temperature, frame_indices in tqdm(
    zip(temperatures, frame_sets), total=len(temperatures)
):
    n_runs = 16 if temperature == "500K" else 7

    base_prefix = f"molecular_dynamics/{temperature}"
    base_destination = f"extracted_data/molecular_dynamics/{temperature}"

    initial_oszicar = Oszicar(f"{base_prefix}/production/0/OSZICAR")
    energies = [step["E0"] for step in initial_oszicar.ionic_steps]
    for idx in range(1, n_runs + 1):
        oszicar = Oszicar(f"{base_prefix}/production/{idx}/OSZICAR")

        energies.extend([step["E0"] for step in oszicar.ionic_steps])

    relaxed_frame_energies = []
    for frame_idx in frame_indices:
        prefix = f"{base_prefix}/selected_frames/{frame_idx}"
        destination = f"{base_destination}/selected_frames/{frame_idx}"
        Path(destination).mkdir(parents=True)

        vasprun = Vasprun(
            f"{prefix}/lobster/vasprun.xml",
            parse_dos=False,
            parse_eigen=False,
            parse_potcar_file=False,
        )
        relaxed_frame_energies.append(vasprun.final_energy)

        magmoms = extract_fe_magmoms(f"{prefix}/lobster/OUTCAR", 36)
        np.save(f"{destination}/magmoms.npy", magmoms)

        shutil.copy(f"{prefix}/lobster/POSCAR", destination)
        shutil.copy(f"{prefix}/lobster/ICOBILIST.lobster", destination)

        for wannier_centres in glob.iglob(f"{prefix}/wannier/localisation/*/*.xyz"):
            shutil.copy(wannier_centres, destination)

    np.save(f"{base_destination}/energies.npy", energies)
    np.save(f"{base_destination}/relaxed_frame_energies.npy", relaxed_frame_energies)

destination = "extracted_data/polymorph_stability"
Path(destination).mkdir()

li2_data = {}
li0_data = {}
print("Extracting polymorph energies and volumes for x = 2, 0")
for composition in tqdm(compositions[::2]):
    for polymorph in polymorphs:
        prefix = f"polymorph_stability/{composition}/{polymorph}"

        vasprun = Vasprun(
            f"{prefix}/vasprun.xml",
            parse_dos=False,
            parse_eigen=False,
            parse_potcar_file=False,
        )

        energy = vasprun.final_energy
        n_atoms = len(vasprun.final_structure)
        volume = vasprun.final_structure.volume

        if composition == "Li2FeSiO4":
            li2_data[polymorph] = {
                "energy": energy,
                "n_atoms": n_atoms,
                "volume": volume,
            }

        else:
            li0_data[polymorph] = {
                "energy": energy,
                "n_atoms": n_atoms,
                "volume": volume,
            }

with open(f"{destination}/li2_data.json", "w") as stream:
    json.dump(li2_data, stream)

with open(f"{destination}/li0_data.json", "w") as stream:
    json.dump(li0_data, stream)

prefix = "hull"
destination = "extracted_data/hull"
Path(destination).mkdir()

stable_configs = ("1x1x1/0", "1x1x1/3", "1x1x1/9")
paths = sorted(glob.glob(f"{prefix}/*/*/vasprun.xml"))
hull_data = []
print("Extracting energies and compositions for the convex hull and voltage curve")
for path in tqdm(paths):
    vasprun = Vasprun(path, parse_dos=False, parse_eigen=False, parse_potcar_file=False)
    energy = vasprun.final_energy

    structure = vasprun.final_structure
    n_atoms = len(structure)
    n_li = len([site for site in structure if site.species_string == "Li"])
    n_o = len([site for site in structure if site.species_string == "O"])

    config_data = {"n_atoms": n_atoms, "n_li": n_li, "n_o": n_o, "energy": energy}

    if any(config in path for config in stable_configs):
        hull_data.insert(0, config_data)

    else:
        hull_data.append(config_data)

li_vasprun = Vasprun(
    f"{prefix}/Li/vasprun.xml",
    parse_dos=False,
    parse_eigen=False,
    parse_potcar_file=False,
)
li_energy = li_vasprun.final_energy

with open(f"{destination}/hull_data.json", "w") as stream:
    json.dump(hull_data, stream)

with open(f"{destination}/li_energy.json", "w") as stream:
    json.dump(li_energy, stream)

prefix = "thermodynamic_stability"
destination = "extracted_data/thermodynamic_stability"
Path(destination).mkdir()

paths = glob.glob(f"{prefix}/*/vasprun.xml")
paths.append("sequential_delithiation/FeSiO4/relaxation/vasprun.xml")
stability_data = {}
print("Extracting thermodynamic stability data for FeSiO4")
for path in tqdm(paths):
    vasprun = Vasprun(path, parse_dos=False, parse_eigen=False, parse_potcar_file=False)
    energy = vasprun.final_energy

    structure = vasprun.final_structure
    n_fu = len(structure) / structure.composition.reduced_composition.num_atoms

    formula = structure.composition.reduced_formula

    stability_data[formula] = {"n_fu": n_fu, "energy": energy}

with open(f"{destination}/stability_data.json", "w") as stream:
    json.dump(stability_data, stream)