"""Define the base solver class."""

from abc import ABC, abstractmethod
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, ClassVar

import numpy as np
from tqdm import tqdm

from ..element import Element, get_element_class
from ..grid import Grid
from ..io import (
    Metrics,
    Observables,
    RecordsBuffer,
    format_profiling_table,
    get_tqdm_file,
    save_results,
)
from ..parameters import InitType, RunParameters
from . import rank, size
from .math_utils import normalize
from .profiling import ProfilingStats, timeit


@dataclass
class BaseSolver(ABC, RunParameters):
    """Abstract data base class for LLG3D solvers."""

    solver_type: ClassVar[str] = "base"  #: Solver type name

    def __post_init__(self) -> None:
        """Initialize the solver after dataclass creation."""
        # Ensure solver name matches class solver_type
        if self.solver != self.solver_type:
            raise ValueError(
                f"Trying to initialize a {self.__class__.__name__}, but solver name  "
                f'mismatch: expected "{self.solver_type}", got "{self.solver}"'
            )

        # Initialize results structure aligned with RunResults
        self.metrics: Metrics = {
            "total_time": 0.0,
            "time_per_ite": 0.0,
            "efficiency": 0.0,
            "CFL": 0.0,
        }
        # Physical observables
        self.observables: Observables = {}
        # Records with optional x_profiles and xyz_average
        self.records: RecordsBuffer = {}

        self.np = size  # Add a parameter for the number of processes

        self.np_float: np.dtype = np.dtype(
            np.float64 if self.precision == "double" else np.float32
        )
        self.grid: Grid = Grid(self.Jx, self.Jy, self.Jz, self.dx)
        if rank == 0:
            print(self.grid)
        # Reference the element class from the element string
        ElementClass: type[Element] = get_element_class(self.element)
        # Pass dtype to Element so its scalar coefficients have the correct
        # numpy dtype (`self.np_float`) and avoid implicit promotion.
        self.elem: Element = ElementClass(
            self.T, self.H_ext, self.grid, self.dt, dtype=self.np_float
        )

        self.rng: np.random.Generator = self._init_rng()
        self.profiling_stats: ProfilingStats = defaultdict(
            lambda: {"time": 0.0, "calls": 0}
        )
        self._tqdm_file = get_tqdm_file()  #: File object for tqdm output

    def theta_init_0(self, t: float) -> np.ndarray:
        """Initialization of theta with 0."""
        return np.zeros(self.grid.dims, dtype=self.np_float)

    def theta_init_dw(self, t: float) -> np.ndarray:
        """Initialization of theta with a domain wall profile."""
        x, _, _ = self.grid.get_mesh(local=size > 1, dtype=self.np_float)
        return 2.0 * np.arctan(
            np.exp(
                -(
                    x
                    - self.grid.Lx / 2
                    + self.elem.d_0 * self.elem.coeff_3 * self.elem.lambda_G * t
                )
                / self.elem.d_0
            )
        )

    def phi_init(self, t: float) -> np.ndarray:
        """Initialization of phi."""
        return (
            np.zeros(self.grid.dims, dtype=self.np_float)
            + self.elem.gamma_0 * self.elem.H_ext * t
        )

    def _init_m_n(self) -> np.ndarray:
        """Initialize the magnetization array at time step n."""
        m_n = np.zeros((3,) + self.grid.dims, dtype=self.np_float)

        if self.init_type == "0":
            theta = self.theta_init_0(0)
        elif self.init_type == "dw":
            theta = self.theta_init_dw(0)
        else:
            raise ValueError(
                f"Unknown initialization type: {self.init_type}, "
                f"should be in {InitType}"
            )

        phi = self.phi_init(0)

        m_n[0] = np.cos(theta)
        m_n[1] = np.sin(theta) * np.cos(phi)
        m_n[2] = np.sin(theta) * np.sin(phi)
        # renormalize to verify the constraint of being on the sphere
        normalize(m_n)
        return m_n

    def _init_rng(self) -> np.random.Generator:
        """
        Initialize a random number generator for temperature fluctuations.

        Returns:
            A numpy random number generator
        """
        # Initialize a sequence of random seeds
        # See: https://numpy.org/doc/stable/reference/random/parallel.html
        ss = np.random.SeedSequence(self.seed)

        # Deploy size x SeedSequence to pass to child processes
        child_seeds = ss.spawn(size)
        streams = [np.random.default_rng(s) for s in child_seeds]
        rng = streams[rank]
        return rng

    @timeit
    def _get_R_random(self) -> np.ndarray:
        """
        Generate the random field for temperature fluctuations.

        Returns:
            Random field array (shape (3, nx, ny, nz))
        """
        R_random = self.elem.coeff_4 * self.rng.standard_normal(
            (3, *self.grid.dims), dtype=self.np_float
        )
        return R_random

    @timeit
    def _normalize(self, m_n: np.ndarray) -> None:
        r"""
        Normalize the magnetization array (in place).

        .. math::

            \mathbf{m}_n = \frac{\mathbf{m}_n}{|\mathbf{m}_n|}

        Args:
            m_n: Magnetization array at time step n (shape (3, nx, ny, nz)).
        """
        normalize(m_n)

    @timeit
    def _xyz_average(self, m: np.ndarray) -> float:
        """
        Returns the spatial average of m with shape (g.dims) using the midpoint method.

        Args:
            m: Array to be integrated

        Returns:
            Spatial average of m
        """
        mm = m.copy()  # copy m to avoid modifying its value

        # on the edges, we divide the contribution by 2
        # x
        mm[0, :, :] /= 2
        mm[-1, :, :] /= 2
        # y
        mm[:, 0, :] /= 2
        mm[:, -1, :] /= 2
        # z
        mm[:, :, 0] /= 2
        mm[:, :, -1] /= 2

        average = mm.sum() / self.grid.ncell
        return float(average)

    def _record_xyz_average(self, m_n: np.ndarray, t: float, n: int) -> None:
        """Update the time average of m1."""
        xyz_average = self._xyz_average(m_n[0])
        if rank == 0:
            # Ensure xyz_average list exists
            if "xyz_average" not in self.records:
                self.records["xyz_average"] = []
            # Record the mean value at time t
            self.records["xyz_average"].append((t, xyz_average))
            # Update time average of m1
            if n >= self.start_averaging:
                # Initialize m1_mean on first use
                if "m1_mean" not in self.observables:
                    self.observables["m1_mean"] = 0.0
                # Accumulate time average (each sample contributes equally)
                self.observables["m1_mean"] += xyz_average

    def _finalize(self) -> None:
        """Normalize m1_mean by the actual number of samples accumulated."""
        if rank == 0:
            if "m1_mean" in self.observables:
                # Divide by actual number of samples
                # (accounting for n_mean sampling interval)
                num_samples = (self.N - self.start_averaging) // self.n_mean
                self.observables["m1_mean"] /= num_samples

    @timeit
    def _yz_average(self, m_i: np.ndarray) -> np.ndarray:
        """
        Returns the spatial average of m using the midpoint method along y and z.

        Args:
            m_i: Array to be integrated

        Returns:
            Spatial average of m in y and z of shape (g.dims[0],)
        """
        # Make a copy of m to avoid modifying its value
        mm = m_i.copy()

        # On y and z edges, divide the contribution by 2
        mm[:, 0, :] /= 2
        mm[:, -1, :] /= 2
        mm[:, :, 0] /= 2
        mm[:, :, -1] /= 2

        n_cell_yz = (mm.shape[1] - 1) * (mm.shape[2] - 1)
        return mm.sum(axis=(1, 2)) / n_cell_yz

    @timeit
    def _update_x_profiles(self, m_n: np.ndarray, t: float) -> None:
        """Update x profiles of the averaged m_i in y and z."""
        # Initialize x_profiles on first use
        if "x_profiles" not in self.records:
            self.records["x_profiles"] = {
                "t": [],
                "m1": [],
                "m2": [],
                "m3": [],
            }
        x_prof = self.records["x_profiles"]
        x_prof["t"].append(t)
        x_prof["m1"].append(self._yz_average(m_n[0]))
        x_prof["m2"].append(self._yz_average(m_n[1]))
        x_prof["m3"].append(self._yz_average(m_n[2]))

    def _record(self, m_n: Any, t: float, n: int) -> None:
        """Record simulation data."""
        # Record the average of m1
        if self.n_mean != 0 and n % self.n_mean == 0:
            self._record_xyz_average(m_n, t, n)
        # Record the x profiles of the averaged m_i in y and z
        if self.n_profile != 0 and n % self.n_profile == 0:
            self._update_x_profiles(m_n, t)

    def _progress_bar(self):
        """Return a progress bar for the given range using tqdm."""
        if self._tqdm_file is None:
            # corresponds to rank != 0 in MPI
            return range(self.N)

        return tqdm(
            range(self.N),
            file=self._tqdm_file,
            dynamic_ncols=True,
            leave=True,
        )

    @abstractmethod
    def _simulate(self) -> float:
        """
        Simulates the system for N iterations.

        Returns:
            total_time: Total simulation time
        """

    def run(self) -> None:
        """Runs the simulation and store the results."""
        total_time = self._simulate()
        self.metrics["total_time"] = total_time
        time_per_ite = total_time / self.N if self.N > 0 else 0.0
        self.metrics["time_per_ite"] = time_per_ite
        self.metrics["efficiency"] = time_per_ite / self.grid.ntot
        self.metrics["CFL"] = float(self.elem.get_CFL())
        if rank == 0:
            # Store only profiling stats for functions/kernels that were actually called
            profiling_filtered = {
                k: v for k, v in self.profiling_stats.items() if v.get("calls", 0) > 0
            }
            if profiling_filtered:
                self.metrics["profiling_stats"] = profiling_filtered

    def _format_profiling(self) -> str:
        """Format the profiling information for display."""
        return format_profiling_table(self.profiling_stats, self.metrics["total_time"])

    def save(self, dir_path: str | Path = ".") -> None:
        """
        Saves the results of the simulation to a .npz file.

        Args:
            dir_path: Directory path to save the results
        """
        if rank == 0:
            if self.metrics["total_time"] > 0:
                s = f"""\
N iterations          = {self.N}
total_time [s]        = {self.metrics["total_time"]:.03f}
time/ite [s/ite]      = {self.metrics["time_per_ite"]:.03e}
efficiency [s/ite/pt] = {self.metrics["efficiency"]:.03e}
CFL                   = {self.metrics["CFL"]:.03e}"""
                print(s)

                # Print profiling info only if enabled
                if self.profiling:
                    print(f"Profiling info:\n{self._format_profiling()}")

            # Export the mean of m1 over space and time
            if "m1_mean" in self.observables:
                print(f"m1_mean           = {self.observables['m1_mean']:e}")

            print(f"Saving {self.result_file}")
            save_results(
                Path(dir_path) / self.result_file,
                self,
                self.metrics,
                observables=self.observables if self.observables else None,
                records_buffer=self.records if self.records else None,
            )
