"""
Post-processes a set of runs.

1. Extracts result data,
2. Plots the computed average magnetization against temperature,
3. Interpolates the computed points using a PCHIP interpolator,
4. Determines the Curie temperature as the value below which the magnetization
   drops under 0.1.
"""

from pathlib import Path


import numpy as np
from scipy.interpolate import PchipInterpolator
from scipy.optimize import brentq

from llg3d.io import load_results


class MagTempData:
    """
    Handle magnetization data using the npz format.

    - Extracts result data,
    - Interpolates the computed points using a PCHIP interpolator,
    - Determines the Curie temperature as the value below which the magnetization
      drops under 0.1.

    Args:
        *files: Paths to the result .npz files
    """

    n_interp = 200  #: number of interpolation points

    def __init__(self, *files: Path | str) -> None:
        #: list of result files
        self.files: list[Path] = [Path(file) for file in files]
        self.params: dict = {}  #: common parameters of the runs
        # Extract data from the runs
        data = self._process_jobs()
        self.temperature = data[:, 0]  #: temperatures from the runs
        self.m1_mean = data[:, 1]  #: mean magnetization
        # Use PCHIP interpolator which preserves positivity and monotonicity
        #: interpolated magnetization function
        self.interp = PchipInterpolator(self.temperature, self.m1_mean)
        self.temperature_interp = np.linspace(
            self.temperature.min(), self.temperature.max(), self.n_interp
        )  #: finer temperature grid for interpolation

    @property
    def T_Curie(self) -> float:
        """
        Return the Curie temperature.

        It is defined as the temperature at which the magnetization equals 0.1,
        found using a root-finding algorithm for precision.

        Returns:
            float: Curie temperature

        Raises:
            ValueError: If the magnetization never crosses 0.1 in the dataset
        """
        # Check if magnetization ever crosses 0.1
        T_min = self.temperature.min()
        T_max = self.temperature.max()
        m1_min = self.interp(T_min)
        m1_max = self.interp(T_max)

        # If 0.1 is never reached
        if (m1_min < 0.1 and m1_max < 0.1) or (m1_min > 0.1 and m1_max > 0.1):
            raise ValueError(
                f"Magnetization never crosses 0.1 in the dataset. "
                f"Range: [{m1_min:.4f}, {m1_max:.4f}]"
            )

        # Find the exact temperature where m = 0.1 using Brent's method
        T_curie = brentq(lambda T: float(self.interp(T)) - 0.1, T_min, T_max)
        return float(T_curie)

    def _process_jobs(self) -> np.ndarray:
        """
        Iterates through calculation directories to assemble data.

        Returns:
            data a numpy array (T, <m>)

        Raises:
            ValueError: If any file does not end with .npz
        """
        for file in self.files:
            if not file.name.endswith(".npz"):
                raise ValueError(f"File {file} should end with .npz")
        # Get parameters from the first run file

        first_results = load_results(self.files[0])
        self.params = first_results.params.as_dict()  # Store common parameters
        data = []
        # Iterating through run directories
        for file in self.files:
            print(f"Processing file: {file}")
            run_results = load_results(file)
            params = run_results.params.as_dict()
            m1_mean = np.nan
            if "observables" in run_results.results:
                m1_mean = run_results.results["observables"].get("m1_mean", np.nan)
            data.append([params["T"], m1_mean])

        data.sort()  # Sorting by increasing temperatures

        return np.array(data)
