src/chemgraph/tools/xanes_core.py

"""Pure-Python XANES/FDMNES helpers (no LangChain / MCP decorators).

Contains all core workflow functions for FDMNES input generation,
execution, result parsing, Materials Project data fetching, and plotting.
Used by the LangChain ``@tool`` wrappers in :mod:`xanes_tools` and the
MCP wrappers in :mod:`chemgraph.mcp.xanes_mcp_parsl`.
"""

from __future__ import annotations

import logging
import os
import pickle
import shutil
import subprocess
from pathlib import Path
from typing import List, Optional

import numpy as np
from ase import Atoms
from ase.io import read as ase_read, write as ase_write

from chemgraph.schemas.xanes_schema import xanes_input_schema, mp_query_schema

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Helper Functions
# ---------------------------------------------------------------------------


def write_fdmnes_input(
    ase_atoms: Atoms,
    z_absorber: int = None,
    input_file_dir: Path = None,
    radius: float = 6.0,
    magnetism: bool = False,
):
    """Write FDMNES input files (fdmfile.txt and fdmnes_in.txt) for a structure.

    Parameters
    ----------
    ase_atoms : ase.Atoms
        Atomic structure to compute XANES for.
    z_absorber : int, optional
        Atomic number of the X-ray absorbing atom.
        Defaults to the heaviest element in the structure.
    input_file_dir : Path, optional
        Directory to write input files into. Defaults to cwd.
    radius : float
        Cluster radius in Angstrom. Default 6.0.
    magnetism : bool
        Enable magnetic contributions. Default False.
    """
    if not isinstance(ase_atoms, Atoms):
        raise TypeError("ase_atoms must be an ase.Atoms object")

    atomic_numbers = ase_atoms.get_atomic_numbers()
    if z_absorber is None:
        z_absorber = int(atomic_numbers.max())

    if input_file_dir is None:
        input_file_dir = Path.cwd()

    with open(input_file_dir / "fdmfile.txt", "w") as f:
        f.write("1\n")
        f.write("fdmnes_in.txt\n")

    with open(input_file_dir / "fdmnes_in.txt", "w") as f:
        f.write("Filout\n")
        f.write(f"{input_file_dir.name}\n\n")

        # Energy mesh
        f.write("Range\n")
        f.write("-55. 1.0 -10. 0.01 5. 0.1 150.\n\n")

        # Cluster radius
        f.write("Radius\n")
        f.write(f"{radius}\n\n")

        # Absorbing atom
        f.write("Z_absorber\n")
        f.write(f"{z_absorber}\n\n")

        # Magnetic contributions
        if magnetism:
            f.write("Magnetism\n\n")

        f.write("Green\n")
        f.write("Density_all\n")
        f.write("Quadrupole\n")
        f.write("Spherical\n")
        f.write("SCF\n\n")

        if all(ase_atoms.pbc):
            f.write("Crystal\n")
            f.write(" ".join(map(str, ase_atoms.cell.cellpar())) + "\n")
            positions = np.round(ase_atoms.get_scaled_positions(), 6)
        else:
            f.write("Molecule\n")
            cell_length = abs(ase_atoms.get_positions().max()) + abs(
                ase_atoms.get_positions().min()
            )
            f.write(f"{cell_length} {cell_length} {cell_length} 90 90 90\n")
            positions = np.round(ase_atoms.get_positions(), 6)

        for i, position in enumerate(positions):
            f.write(f"{atomic_numbers[i]} " + " ".join(map(str, position)) + "\n")

        f.write("\n")
        f.write("Convolution\n")
        f.write("End")


def get_normalized_xanes(
    conv_file: Path | str,
    pre_edge_width: float = 20.0,
    post_edge_width: float = 50.0,
    calc_E0: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
    """Normalize a XANES spectrum from an FDMNES convolution output file.

    Parameters
    ----------
    conv_file : Path or str
        Path to the FDMNES ``*_conv.txt`` output file.
    pre_edge_width : float
        Width of the pre-edge region in eV for baseline fitting.
    post_edge_width : float
        Width of the post-edge region in eV for step normalization.
    calc_E0 : bool
        If True, determine the edge energy E0 from the maximum of dmu/dE.
        Otherwise E0 is assumed to be 0 (the FDMNES convention).

    Returns
    -------
    normalized : np.ndarray
        (N, 2) array of [energy, normalized_mu].
    raw : np.ndarray
        (N, 2) array of [energy, raw_mu] as read from the file.
    """
    energy_xas = np.loadtxt(conv_file, skiprows=1)

    E = energy_xas[:, 0].astype(float)
    mu = energy_xas[:, 1].astype(float)

    if calc_E0:
        dmu_dE = np.gradient(mu, E)
        E0 = E[np.argmax(dmu_dE)]
    else:
        E0 = 0

    pre_mask = E <= (E0 - pre_edge_width)
    post_mask = E >= (E0 + post_edge_width)

    m_pre, b_pre = np.polyfit(E[pre_mask], mu[pre_mask], 1)
    m_post, b_post = np.polyfit(E[post_mask], mu[post_mask], 1)

    pre_line = m_pre * E + b_pre
    mu_corr = mu - pre_line

    step = (m_post * E0 + b_post) - (m_pre * E0 + b_pre)
    mu_norm = mu_corr / step

    return np.column_stack([E, mu_norm]), energy_xas


def extract_conv(fdmnes_output_dir: Path | str) -> dict:
    """Extract all convolution output files from an FDMNES run directory.

    Parameters
    ----------
    fdmnes_output_dir : Path or str
        Directory containing FDMNES output files.

    Returns
    -------
    dict
        Mapping of index to (N, 2) arrays of [energy, mu].
    """
    if not isinstance(fdmnes_output_dir, Path):
        fdmnes_output_dir = Path(fdmnes_output_dir)

    energy_xas = {}
    for i, conv_file in enumerate(fdmnes_output_dir.glob("*conv.txt")):
        energy_xas[i] = np.loadtxt(conv_file, skiprows=1)

    return energy_xas


# ---------------------------------------------------------------------------
# Data directory helper
# ---------------------------------------------------------------------------


def _get_data_dir() -> Path:
    """Return the working data directory for XANES workflows."""
    cwd = Path.cwd()
    if "PBS_O_WORKDIR" in os.environ:
        cwd = Path(os.environ["PBS_O_WORKDIR"])

    data_dir = cwd / "xanes_data"
    if not data_dir.exists():
        data_dir.mkdir(parents=True)
    return data_dir


# ---------------------------------------------------------------------------
# Core Workflow Functions
# ---------------------------------------------------------------------------


def run_xanes_core(params: xanes_input_schema) -> dict:
    """Run a single XANES/FDMNES calculation for one structure.

    This is the core function analogous to ``run_graspa_core``. It:
    1. Reads the input structure file via ASE.
    2. Creates FDMNES input files via ``write_fdmnes_input``.
    3. Runs FDMNES via subprocess.
    4. Parses the convolution output if available.

    Parameters
    ----------
    params : xanes_input_schema
        Input parameters for the FDMNES calculation.

    Returns
    -------
    dict
        Result dictionary with keys: status, output_dir, conv_data (if success),
        error (if failure).
    """
    fdmnes_exe = os.environ.get("FDMNES_EXE")
    if not fdmnes_exe:
        raise ValueError(
            "FDMNES_EXE environment variable is not set. "
            "Set it to the path of the FDMNES executable."
        )

    input_path = Path(params.input_structure_file).resolve()
    if not input_path.exists():
        raise FileNotFoundError(f"Input structure file not found: {input_path}")

    atoms = ase_read(str(input_path))

    # Determine output directory
    if params.output_dir is not None:
        run_dir = Path(params.output_dir).resolve()
    else:
        run_dir = input_path.parent / f"fdmnes_{input_path.stem}"
    run_dir.mkdir(parents=True, exist_ok=True)

    # Write FDMNES input files
    write_fdmnes_input(
        ase_atoms=atoms,
        z_absorber=params.z_absorber,
        input_file_dir=run_dir,
        radius=params.radius,
        magnetism=params.magnetism,
    )

    # Save the atoms object alongside the inputs for provenance
    formula = atoms.get_chemical_formula()
    z_abs = params.z_absorber or int(atoms.get_atomic_numbers().max())
    mp_id = atoms.info.get("MP-id", "local")
    pkl_filename = f"Z{z_abs}_{mp_id}_{formula}.pkl"
    with open(run_dir / pkl_filename, "wb") as f:
        pickle.dump(atoms, f)

    # Run FDMNES
    logger.info("Running FDMNES in %s", run_dir)
    with (
        open(run_dir / "fdmnes_stdout.txt", "w") as fp_out,
        open(run_dir / "fdmnes_stderr.txt", "w") as fp_err,
    ):
        proc = subprocess.run(
            fdmnes_exe,
            cwd=str(run_dir),
            stdout=fp_out,
            stderr=fp_err,
            shell=True,
        )

    if proc.returncode != 0:
        logger.error(
            "FDMNES failed with return code %d in %s", proc.returncode, run_dir
        )
        return {
            "status": "failure",
            "output_dir": str(run_dir),
            "error": f"FDMNES exited with return code {proc.returncode}",
        }

    # Parse results
    conv_data = extract_conv(run_dir)
    if not conv_data:
        logger.warning("No convolution output found in %s", run_dir)
        return {
            "status": "failure",
            "output_dir": str(run_dir),
            "error": "No *conv.txt output files found after FDMNES execution.",
        }

    logger.info("FDMNES completed successfully in %s", run_dir)
    return {
        "status": "success",
        "output_dir": str(run_dir),
        "n_conv_files": len(conv_data),
    }


def fetch_materials_project_data(
    params: mp_query_schema,
    db_path: Path,
) -> dict:
    """Fetch optimized structures from Materials Project.

    Parameters
    ----------
    params : mp_query_schema
        Query parameters including chemical formulas and API key.
    db_path : Path
        Directory to save the fetched structures.

    Returns
    -------
    dict
        atoms_list : list[Atoms]    -- fetched ASE Atoms objects
        structure_files : list[str] -- absolute paths to saved CIF files
        pickle_file : str           -- absolute path to atoms_db.pkl
        n_structures : int          -- number of structures fetched
    """
    from mp_api.client import MPRester
    from pymatgen.io.ase import AseAtomsAdaptor

    api_key = params.mp_api_key or os.environ.get("MP_API_KEY")
    if not api_key:
        raise ValueError(
            "No Materials Project API key provided. "
            "Pass it via mp_api_key or set the MP_API_KEY environment variable."
        )

    logger.info("Fetching data from Materials Project for: %s", params.chemsys)
    atoms_list = []

    with MPRester(api_key) as mpr:
        doc_list = mpr.materials.summary.search(
            fields=["material_id", "structure"],
            energy_above_hull=(0, params.energy_above_hull),
            formula=params.chemsys,
            deprecated=False,
        )

        for doc in doc_list:
            ase_atoms = AseAtomsAdaptor.get_atoms(doc.structure)
            ase_atoms.info.update({"MP-id": str(doc.material_id)})
            atoms_list.append(ase_atoms)

    if not db_path.exists():
        db_path.mkdir(parents=True)

    # Save pickle database
    pkl_path = db_path / "atoms_db.pkl"
    with open(pkl_path, "wb") as f:
        pickle.dump(atoms_list, f)

    # Save individual CIF files
    structure_files = []
    for atoms in atoms_list:
        mp_id = atoms.info.get("MP-id", "unknown")
        formula = atoms.get_chemical_formula()
        cif_path = db_path / f"{mp_id}_{formula}.cif"
        ase_write(str(cif_path), atoms)
        structure_files.append(str(cif_path))

    logger.info(
        "Saved %d structures (%s) and pickle database to %s",
        len(atoms_list),
        [Path(f).name for f in structure_files],
        db_path,
    )

    return {
        "atoms_list": atoms_list,
        "structure_files": structure_files,
        "pickle_file": str(pkl_path),
        "n_structures": len(atoms_list),
    }


def create_fdmnes_inputs(
    root_dir: Path,
    atoms_list: Optional[List[Atoms]] = None,
    z_absorber: Optional[int] = None,
    radius: float = 6.0,
    magnetism: bool = False,
) -> Path:
    """Create FDMNES input files for a batch of structures.

    Parameters
    ----------
    root_dir : Path
        Root directory for the batch. A ``fdmnes_batch_runs`` subdirectory
        will be created containing per-structure run directories.
    atoms_list : list[ase.Atoms], optional
        Structures to process. If None, loads from ``root_dir/atoms_db.pkl``.
    z_absorber : int, optional
        Atomic number of the absorbing atom. Defaults to heaviest per structure.
    radius : float
        Cluster radius in Angstrom.
    magnetism : bool
        Enable magnetic contributions.

    Returns
    -------
    Path
        Path to the ``fdmnes_batch_runs`` directory.
    """
    logger.info("Creating FDMNES inputs in %s", root_dir)
    runs_dir = root_dir / "fdmnes_batch_runs"

    start_idx = 0
    if runs_dir.exists():
        for subdir in runs_dir.iterdir():
            try:
                start_idx = max(start_idx, int(subdir.name.split("_")[-1]))
            except ValueError:
                continue
        last_run = runs_dir / f"run_{start_idx}"
        if last_run.exists():
            shutil.rmtree(last_run)
    else:
        runs_dir.mkdir(parents=True)

    if atoms_list is None:
        db_path = root_dir / "atoms_db.pkl"
        if not db_path.exists():
            raise FileNotFoundError(f"No atoms provided and {db_path} not found.")
        with open(db_path, "rb") as f:
            atoms_list = pickle.load(f)

    for i, atoms in enumerate(atoms_list, start=start_idx):
        curr_run_dir = runs_dir / f"run_{i}"
        curr_run_dir.mkdir(parents=True, exist_ok=True)

        current_z = (
            z_absorber
            if z_absorber is not None
            else int(max(atoms.get_atomic_numbers()))
        )
        write_fdmnes_input(
            ase_atoms=atoms,
            input_file_dir=curr_run_dir,
            z_absorber=current_z,
            radius=radius,
            magnetism=magnetism,
        )

        mp_id = atoms.info.get("MP-id", "local")
        formula = atoms.get_chemical_formula()
        pkl_filename = f"Z{current_z}_{mp_id}_{formula}.pkl"
        with open(curr_run_dir / pkl_filename, "wb") as f:
            pickle.dump(atoms, f)

    return runs_dir


def expand_database_results(root_dir: Path, runs_dir: Path) -> None:
    """Expand the atoms database with XANES convolution results.

    For each completed run directory, loads the pickled Atoms object,
    attaches the FDMNES convolution data to ``atoms.info``, and saves
    all expanded structures to ``root_dir/atoms_db_expanded.pkl``.

    Parameters
    ----------
    root_dir : Path
        Root directory where the expanded database will be saved.
    runs_dir : Path
        Directory containing ``run_*`` subdirectories with FDMNES outputs.
    """
    logger.info("Expanding database with XANES results...")
    expanded_atoms_list = []

    for sub_dir in sorted(runs_dir.glob("run_*")):
        atoms_pkl_files = list(sub_dir.glob("*.pkl"))
        if not atoms_pkl_files:
            continue

        with open(atoms_pkl_files[0], "rb") as f:
            ase_atoms = pickle.load(f)

        conv_data = extract_conv(fdmnes_output_dir=sub_dir)
        ase_atoms.info.update({"FDMNES-xanes": conv_data})
        expanded_atoms_list.append(ase_atoms)

    with open(root_dir / "atoms_db_expanded.pkl", "wb") as f:
        pickle.dump(expanded_atoms_list, f)

    logger.info(
        "Saved %d expanded structures to %s",
        len(expanded_atoms_list),
        root_dir / "atoms_db_expanded.pkl",
    )


def plot_xanes_results(root_dir: Path, runs_dir: Path) -> dict:
    """Generate normalized XANES plots for completed FDMNES calculations.

    For each run directory containing a ``*_conv.txt`` file, produces
    a ``xanes_plot.png`` with the normalized absorption spectrum.

    Parameters
    ----------
    root_dir : Path
        Root data directory (unused currently, reserved for summary plots).
    runs_dir : Path
        Directory containing ``run_*`` subdirectories with FDMNES outputs.

    Returns
    -------
    dict
        plot_files : list[str]  -- absolute paths to generated plot images
        n_plots : int           -- number of plots successfully generated
        n_failed : int          -- number of runs that failed to plot
        failed : list[str]      -- names of run directories that failed
    """
    import matplotlib.pyplot as plt

    logger.info("Plotting XANES results from %s", runs_dir)

    plot_files = []
    failed = []

    for sub_dir in sorted(runs_dir.glob("run_*")):
        conv_file = next(sub_dir.glob("*_conv.txt"), None)
        if conv_file:
            try:
                norm_energy, _raw = get_normalized_xanes(conv_file)
                plot_path = sub_dir / "xanes_plot.png"
                plt.figure()
                plt.plot(norm_energy[:, 0], norm_energy[:, 1], label=sub_dir.name)
                plt.xlabel("Energy [eV]")
                plt.ylabel("Normalized Absorption")
                plt.title(f"XANES for {sub_dir.name}")
                plt.legend()
                plt.savefig(plot_path, dpi=150)
                plt.close()
                plot_files.append(str(plot_path))
                logger.info("Plotted %s", sub_dir.name)
            except Exception as e:
                logger.error("Failed to plot %s: %s", sub_dir.name, e)
                failed.append(sub_dir.name)

    return {
        "plot_files": plot_files,
        "n_plots": len(plot_files),
        "n_failed": len(failed),
        "failed": failed,
    }