"""
The primary xarray accessor namespace for the xmris package.

This module exposes the `.xmr` namespace to xarray DataArrays and Datasets.
It uses a "Hybrid Mixin" pattern: the user-facing API remains perfectly flat
for fluent method chaining (e.g., `da.xmr.apodize_exp().xmr.fft()`), while
the underlying developer API is strictly modularized into Mixin classes.
"""

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import xarray as xr

# Import our core architecture
from xmris.core.config import ATTRS, COORDS, DIMS
from xmris.core.utils import _check_dims, as_variable
from xmris.core.validation import requires_attrs

# Processing imports
from xmris.processing.fid import apodize_exp, apodize_lg, to_fid, to_spectrum, zero_fill
from xmris.processing.fourier import fft, fftc, fftshift, ifft, ifftc, ifftshift
from xmris.processing.phasing import autophase, phase

# =============================================================================
# Sub-Accessors (Terminal / Visualization tools)
# =============================================================================
from xmris.vendor.bruker import remove_digital_filter

# Deferred plot configs
from xmris.visualization.plot import PlotHeatmapConfig, PlotRidgeConfig


class XmrisDatasetPlotAccessor:
    """Sub-accessor for xmris xr.Datasets plotting functionalities."""

    def __init__(self, obj: xr.Dataset):
        self._obj = obj

    def trajectory(
        self,
        dim: str,
        metabolites: list[str] | None = None,
        ax: plt.Axes | None = None,
        config=None,
    ):
        """Plot kinetic trajectories with CRLB shading."""
        from xmris.visualization.plot.plot_trajectory import plot_trajectory

        return plot_trajectory(
            self._obj, dim=dim, metabolites=metabolites, ax=ax, config=config
        )

    def qc_grid(self, dim: str, config=None):
        """Plot a grid of spectra and fits to quickly visually inspect quality."""
        from xmris.visualization.plot.plot_qc_grid import plot_qc_grid

        return plot_qc_grid(self._obj, dim=dim, config=config)


class XmrisPlotAccessor:
    """Sub-accessor for xmris plotting functionalities (accessed via .xmr.plot)."""

    def __init__(self, obj: xr.DataArray):
        self._obj = obj

    def ridge(
        self,
        x_dim: str | None = None,
        stack_dim: str | None = None,
        ax: plt.Axes | None = None,
        config: PlotRidgeConfig | None = None,
    ) -> plt.Axes:
        """Generate a ridge plot (2D waterfall) of stacked 1D spectra."""
        from xmris.visualization.plot import plot_ridge as _plot_ridge

        return _plot_ridge(
            da=self._obj, x_dim=x_dim, stack_dim=stack_dim, ax=ax, config=config
        )

    def heatmap(
        self,
        x_dim: str | None = None,
        stack_dim: str | None = None,
        ax: plt.Axes | None = None,
        config: PlotHeatmapConfig | None = None,
    ) -> plt.Axes:
        """Generate a 2D heatmap plot of stacked 1D spectra."""
        from xmris.visualization.plot.plot_heatmap import plot_heatmap as _plot_heatmap

        return _plot_heatmap(
            da=self._obj, x_dim=x_dim, stack_dim=stack_dim, ax=ax, config=config
        )


class XmrisWidgetAccessor:
    """Sub-accessor for xmris interactive widget functionalities.

    This class provides a dedicated namespace for interactive UI components
    powered by AnyWidget. It is accessed via the `.xmr.widget` attribute
    on an xarray DataArray.
    """

    def __init__(self, obj: xr.DataArray):
        """
        Initialize the widget sub-accessor.

        Parameters
        ----------
        obj : xr.DataArray
            The underlying xarray DataArray object being operated on.
        """
        self._obj = obj

    def phase_spectrum(
        self,
        width: int = 740,
        height: int = 400,
        show_grid: bool = True,
        show_pivot: bool = True,
        **kwargs,
    ):
        """Open an interactive zero- and first-order phase correction widget.

        This method launches an AnyWidget-based user interface directly in the
        Jupyter Notebook. It allows for manual, real-time adjustment of the
        zero-order (p0) and first-order (p1) phase angles of a 1-D
        complex-valued NMR/MRS spectrum.

        Parameters
        ----------
        width : int, optional
            Width of the widget in pixels. Default is 740.
        height : int, optional
            Height of the widget in pixels. Default is 400.
        show_grid : bool, optional
            Toggle the background grid visibility. Default is True.
        show_pivot : bool, optional
            Toggle the visibility of the p1 pivot indicator. Default is True.
        **kwargs
            Additional arguments passed to the underlying PhaseWidget.

        Returns
        -------
        PhaseWidget
            The interactive widget instance. Assigning this to a variable allows
            you to programmatically extract the optimized phase angles after
            interacting with the UI.

        Raises
        ------
        ValueError
            If the underlying DataArray is not 1-dimensional or does not contain
            complex-valued data.

        Notes
        -----
        - **Zero-order phase (p0)**: Adjusts the phase uniformly across the spectrum.
        - **First-order phase (p1)**: Adjusts phase linearly relative to a pivot point.
        - The pivot point (p_pivot) is automatically set to the coordinate
          corresponding to the maximum magnitude peak.

        """
        # Lazy import to avoid loading AnyWidget/frontend assets unless requested
        from xmris.visualization.widget import phase_spectrum

        # The underlying function handles the 1-D and complex-type validation
        return phase_spectrum(
            self._obj,
            width=width,
            height=height,
            show_grid=show_grid,
            show_pivot=show_pivot,
            **kwargs,
        )


# =============================================================================
# Mixins (Developer API Modularity)
# =============================================================================


class XmrisSpectrumCoordsMixin:
    """Mixin providing operations to translate physical coordinate systems."""

    @requires_attrs(ATTRS.reference_frequency, ATTRS.carrier_ppm)
    def to_ppm(self, dim: str = DIMS.frequency) -> xr.DataArray:
        """Convert relative frequency axis [Hz] to absolute chemical shift axis [ppm]."""
        _check_dims(self._obj, dim, "to_ppm")

        mhz = self._obj.attrs[ATTRS.reference_frequency]
        carrier_ppm = self._obj.attrs[ATTRS.carrier_ppm]
        hz_coords = self._obj.coords[dim].values

        # 1. Calculate the math
        ppm_coords = carrier_ppm + (hz_coords / mhz)

        # 2. Build the fully-formed xarray Variable (data + metadata)
        shift_var = as_variable(DIMS.chemical_shift, dim, ppm_coords)

        # 3. Assign and swap in one clean sweep
        obj = self._obj.assign_coords({DIMS.chemical_shift: shift_var})
        return obj.swap_dims({dim: DIMS.chemical_shift})

    @requires_attrs(ATTRS.reference_frequency, ATTRS.carrier_ppm)
    def to_hz(self, dim: str = DIMS.chemical_shift) -> xr.DataArray:
        """Convert absolute chemical shift axis [ppm] to relative frequency axis [Hz]."""
        _check_dims(self._obj, dim, "to_hz")

        mhz = self._obj.attrs[ATTRS.reference_frequency]
        carrier_ppm = self._obj.attrs[ATTRS.carrier_ppm]
        ppm_coords = self._obj.coords[dim].values

        hz_coords = (ppm_coords - carrier_ppm) * mhz

        # Pack the data and metadata together instantly
        freq_var = as_variable(COORDS.frequency, dim, hz_coords)

        obj = self._obj.assign_coords({COORDS.frequency: freq_var})
        return obj.swap_dims({dim: DIMS.frequency})


class XmrisFourierMixin:
    """Mixin providing generalized N-dimensional Fourier transforms and shifts."""

    def fftshift(self, dim: str | list[str]) -> xr.DataArray:
        """
        Apply fftshift by rolling data and coordinates along specified dimensions.

        Moves the zero-frequency component to the center of the spectrum.
        """
        return fftshift(self._obj, dim=dim)

    def ifftshift(self, dim: str | list[str]) -> xr.DataArray:
        """
        Apply ifftshift by rolling data and coordinates along specified dimensions.

        The exact inverse of :meth:`fftshift`.
        """
        return ifftshift(self._obj, dim=dim)

    def fft(
        self,
        dim: str | list[str] = DIMS.time,
        out_dim: str | list[str] | None = None,
    ) -> xr.DataArray:
        """
        Perform a standard N-dimensional Fast Fourier Transform (no shifts).

        Parameters
        ----------
        dim : str or list of str, optional
            Dimension(s) to transform, by default `DIMS.time`.
        out_dim : str or list of str, optional
            Optional new dimension name(s), by default None.

        Returns
        -------
        xr.DataArray
            The transformed DataArray.
        """
        return fft(self._obj, dim=dim, out_dim=out_dim)

    def ifft(
        self,
        dim: str | list[str] = DIMS.frequency,
        out_dim: str | list[str] | None = None,
    ) -> xr.DataArray:
        """
        Perform a standard N-dimensional Inverse FFT (no shifts).

        Parameters
        ----------
        dim : str or list of str, optional
            Dimension(s) to transform, by default `DIMS.frequency`.
        out_dim : str or list of str, optional
            Optional new dimension name(s), by default None.

        Returns
        -------
        xr.DataArray
            The transformed DataArray.
        """
        return ifft(self._obj, dim=dim, out_dim=out_dim)

    def fftc(
        self,
        dim: str | list[str] = DIMS.time,
        out_dim: str | list[str] | None = None,
    ) -> xr.DataArray:
        """Perform a centered N-dimensional FFT (ifftshift -> fft -> fftshift)."""
        return fftc(self._obj, dim=dim, out_dim=out_dim)

    def ifftc(
        self,
        dim: str | list[str] = DIMS.frequency,
        out_dim: str | list[str] | None = None,
    ) -> xr.DataArray:
        """Perform a centered N-dimensional Inverse FFT (ifftshift -> ifft -> fftshift)."""  # noqa: E501
        return ifftc(self._obj, dim=dim, out_dim=out_dim)


class XmrisProcessingMixin:
    """Mixin providing common NMR/MRI Free Induction Decay processing tools."""

    def apodize_exp(self, dim: str = DIMS.time, lb: float = 1.0) -> xr.DataArray:
        """
        Multiply the time-domain signal by a decreasing mono-exponential filter.

        Parameters
        ----------
        dim : str, optional
            The dimension corresponding to time, by default `DIMS.time`.
        lb : float, optional
            The desired line broadening factor in Hz, by default 1.0.

        Returns
        -------
        xr.DataArray
            A new apodized DataArray, preserving coordinates and attributes.
        """
        return apodize_exp(self._obj, dim=dim, lb=lb)

    def apodize_lg(
        self, dim: str = DIMS.time, lb: float = 1.0, gb: float = 1.0
    ) -> xr.DataArray:
        """
        Apply a Lorentzian-to-Gaussian transformation filter.

        Parameters
        ----------
        dim : str, optional
            The dimension corresponding to time, by default `DIMS.time`.
        lb : float, optional
            The Lorentzian line broadening to cancel in Hz, by default 1.0.
        gb : float, optional
            The Gaussian line broadening to apply in Hz, by default 1.0.

        Returns
        -------
        xr.DataArray
            A new apodized DataArray, preserving coordinates and attributes.
        """
        return apodize_lg(self._obj, dim=dim, lb=lb, gb=gb)

    def to_spectrum(
        self, dim: str = DIMS.time, out_dim: str = DIMS.frequency
    ) -> xr.DataArray:
        """
        Convert a time-domain FID to a frequency-domain spectrum.

        Parameters
        ----------
        dim : str, optional
            The time dimension to transform, by default `DIMS.time`.
        out_dim : str, optional
            The name of the resulting frequency dimension, by default `DIMS.frequency`.

        Returns
        -------
        xr.DataArray
            The centered frequency-domain spectrum.
        """
        return to_spectrum(self._obj, dim=dim, out_dim=out_dim)

    def to_fid(self, dim: str = DIMS.frequency, out_dim: str = DIMS.time) -> xr.DataArray:
        """
        Convert a frequency-domain spectrum to a time-domain FID.

        Parameters
        ----------
        dim : str, optional
            The frequency dimension to transform, by default `DIMS.frequency`.
        out_dim : str, optional
            The name of the resulting time dimension, by default `DIMS.time`.

        Returns
        -------
        xr.DataArray
            The un-shifted time-domain FID data.
        """
        return to_fid(self._obj, dim=dim, out_dim=out_dim)

    def zero_fill(
        self,
        dim: str = DIMS.time,
        target_points: int = 1024,
        position: str = "end",
    ) -> xr.DataArray:
        """
        Pad the specified dimension with zero amplitude points.

        Parameters
        ----------
        dim : str, optional
            The dimension along which to pad zeros, by default `DIMS.time`.
        target_points : int, optional
            The total number of points desired after padding, by default 1024.
        position : {"end", "symmetric"}, optional
            Where to apply the zeros. Use "end" for time-domain FIDs, and
            "symmetric" for spatial frequency domains like k-space. By default "end".

        Returns
        -------
        xr.DataArray
            A new DataArray padded with zeros to the target length.
        """
        return zero_fill(
            self._obj, dim=dim, target_points=target_points, position=position
        )


class XmrisPhasingMixin:
    """Mixin providing common MR spectra phasing tools."""

    def phase(
        self,
        dim: str = DIMS.frequency,
        p0: float = 0.0,
        p1: float = 0.0,
        pivot: float = None,
    ) -> xr.DataArray:
        """
        Apply zero- and first-order phase correction to the spectrum.

        Parameters
        ----------
        dim : str, optional
            The frequency dimension along which to apply phase correction,
            by default `DIMS.frequency`.
        p0 : float, optional
            Zero-order phase angle in degrees, by default 0.0.
        p1 : float, optional
            First-order phase angle in degrees, by default 0.0.
        pivot : float, optional
            The coordinate value (e.g., ppm or Hz) around which p1 is pivoted.
            If None, standard nmrglue index-0 pivoting is used.

        Returns
        -------
        xr.DataArray
            The phase-corrected spectrum with phase_p0 and phase_p1 stored
            in the attributes.
        """
        return phase(self._obj, dim=dim, p0=p0, p1=p1, pivot=pivot)

    def autophase(
        self,
        dim: str = DIMS.frequency,
        method: str = "acme",
        peak_width: int = 100,
        lb: float = 0.0,
        temp_time_dim: str = DIMS.time,
        **kwargs,
    ) -> xr.DataArray:
        """
        Automatically calculate and apply phase correction to a spectrum.

        Parameters
        ----------
        da : xr.DataArray
            The input frequency-domain spectrum.
        dim : str, optional
            The coordinate dimension to operate on, by default `DIMS.frequency`.
        method : {"acme", "peak_minima", "positivity"}, optional
            The scoring algorithm to use. "acme" relies on entropy and is best for
            multi-peak high SNR spectra. "positivity" and "peak_minima" are optimized
            for sparse/noisy spectra. By default "acme".
        peak_width : float, optional
            Width of the ROI (in units of `dim`, e.g., Hz or ppm) for the local methods.
            Concentrates the solver on the region surrounding the target peak.
            By default 0.5.
        target_coord : float | None, optional
            The explicit coordinate (e.g. 171.0 ppm) to target for local methods.
            If None, the coordinate of the maximum absolute magnitude is used.
        p0_only : bool, optional
            If True, locks p1=0 and only optimizes the zero-order phase. Highly
            recommended for sparse spectra evaluated over a narrow `peak_width`.
        lb : float, optional
            Optional exponential line broadening (in Hz). Can help smooth extreme
            noise for ACME, but usually unnecessary for local methods. By default 0.0.
        temp_time_dim : str, optional
            The name used for the temporary time dimension if lb > 0.
        **kwargs :
            Additional keyword arguments passed to `scipy.optimize.differential_evolution`.

        Returns
        -------
        xr.DataArray
            The phased spectrum.
        """  # noqa: E501
        return autophase(
            self._obj,
            dim=dim,
            method=method,
            peak_width=peak_width,
            lb=lb,
            temp_time_dim=temp_time_dim,
            **kwargs,
        )


# =============================================================================
# Main User API Registration
# =============================================================================


@xr.register_dataset_accessor("xmr")
class XmrisDatasetAccessor:
    """Accessor for xmris xr.Datasets (e.g., fitting results)."""

    def __init__(self, xarray_obj: xr.Dataset):
        self._obj = xarray_obj
        self._plot = None

    @property
    def plot(self) -> XmrisDatasetPlotAccessor:
        """Access xmris plotting functionalities."""
        if self._plot is None:
            self._plot = XmrisDatasetPlotAccessor(self._obj)
        return self._plot


@xr.register_dataarray_accessor("xmr")
class XmrisAccessor(
    XmrisSpectrumCoordsMixin, XmrisFourierMixin, XmrisProcessingMixin, XmrisPhasingMixin
):
    """
    Main Accessor for xarray DataArrays to perform MRI and MRS operations.

    This class is registered under the `.xmr` namespace. It inherits from
    several domain-specific Mixins to provide a fluent, method-chaining API
    (e.g., `da.xmr.apodize_exp().xmr.to_spectrum().xmr.to_ppm()`) without
    creating an unmanageable monolithic class.

    Attributes
    ----------
    _obj : xr.DataArray
        The underlying xarray DataArray object being operated on.
    """

    def __init__(self, xarray_obj: xr.DataArray):
        """Initialize the accessor with the xarray object."""
        self._obj = xarray_obj
        self._plot = None  # Cache for the plot sub-accessor
        self._widget = None  # Cache for the widget sub-accessor

    @property
    def plot(self) -> XmrisPlotAccessor:
        """Access xmris plotting functionalities for DataArrays."""
        if self._plot is None:
            self._plot = XmrisPlotAccessor(self._obj)
        return self._plot

    @property
    def widget(self) -> XmrisWidgetAccessor:
        """Access xmris plotting functionalities for DataArrays."""
        if self._widget is None:
            self._widget = XmrisWidgetAccessor(self._obj)
        return self._widget

    # --- Fitting ---

    def fit_amares(
        self,
        prior_knowledge_file: str | Path,
        dim: str = "time",
        mhz: float | None = None,
        sw: float | None = None,
        deadtime: float | None = None,
        method: str = "leastsq",
        initialize_with_lm: bool = True,
        num_workers: int = 4,
        init_fid: np.ndarray | None = None,
        **kwargs,
    ) -> xr.Dataset:
        """
        Apply AMARES time-domain fitting to an N-dimensional FID.

        This method wraps `pyAMARES` to perform parallelized batch fitting
        across spatial or repetition dimensions. The numerical results and
        the reconstructed time-domain fits are packed into an aligned xarray Dataset.

        Requires the optional `pyAMARES` package to be installed.

        Parameters
        ----------
        prior_knowledge_file : str | Path
            Path to the CSV or XLSX file containing the prior knowledge constraints.
        dim : str, optional
            The time dimension along which to fit, by default "time".
        mhz : float, optional
            Spectrometer frequency in MHz. If None, attempts to read from attrs['MHz'].
        sw : float, optional
            Spectral width in Hz. If None, attempts to calculate from `dim` coordinates.
        deadtime : float, optional
            Time delay before the first point in seconds. If None, defaults to 0.0.
        method : {"leastsq", "least_squares"}, optional
            Fitting method. Defaults to 'leastsq' (Levenberg-Marquardt).
        initialize_with_lm : bool, optional
            Run an internal Levenberg-Marquardt initializer before fitting.
            Defaults to True.
        num_workers : int, optional
            Number of parallel processes to spawn. Defaults to 4.
        init_fid : np.ndarray, optional
            A 1D complex array to use as the template for pyAMARES initialization.
            If None, the function automatically selects the spectrum with the highest SNR.

        Returns
        -------
        xr.Dataset
            A dataset containing the original data, the fitted FIDs, the residuals,
            and the quantified parameters (amplitude, chem_shift, linewidth, phase,
            CRLB, SNR) mapped across the original dimensions and the new 'Metabolite'
            dimension.

        Raises
        ------
        ImportError
            If the `pyAMARES` package is not installed.
        """
        try:
            from xmris.fitting.amares import fit_amares as _internal_fit_amares
        except ImportError as e:
            raise ImportError(
                "The '.fit_amares()' method requires the optional 'pyAMARES' package. "
                "Please install it using 'pip install pyAMARES' or 'uv add pyAMARES'."
            ) from e

        return _internal_fit_amares(
            self._obj,
            prior_knowledge_file=prior_knowledge_file,
            dim=dim,
            mhz=mhz,
            sw=sw,
            deadtime=deadtime,
            method=method,
            initialize_with_lm=initialize_with_lm,
            num_workers=num_workers,
            init_fid=init_fid,
            **kwargs,
        )

    # --- Vendor Specific ---

    def remove_digital_filter(
        self, group_delay: float, dim: str = "time", keep_length: bool = True
    ) -> xr.DataArray:
        """
        Remove the hardware digital filter group delay from Bruker FID data.

        Bruker consoles use a cascade of digital FIR filters during analog-to-digital
        conversion. Because these filters calculate a moving average, they require time
        to "wake up", introducing a causality delay at the start of the Free Induction
        Decay (FID). This manifests as a time-shift, effectively prepending the actual
        signal with a specific number of filter transient points.

        Parameters
        ----------
        group_delay : float
            The exact delay value to remove. This should be read directly from the
            Bruker `ACQ_RxFilterInfo` parameter array.
        dim : str, optional
            The time dimension along which to apply the correction, by default "time".
        keep_length : bool, optional
            If True, appends pure zeros to the end of the FID to replace the truncated
            startup points, maintaining the original length. By default True.

        Returns
        -------
        xr.DataArray
            The corrected FID data with the filter transient stripped and phase aligned.
        """
        return remove_digital_filter(
            self._obj, group_delay=group_delay, dim=dim, keep_length=keep_length
        )

    # --- Utility / Formatting ---

    def to_real_imag(
        self, dim: str = DIMS.component, coords: tuple[str, str] = ("real", "imag")
    ) -> xr.DataArray:
        """Split a complex array into a real-valued array with an extra component dimension."""  # noqa: E501
        # Kept the deferred import to match your original package load-time strategy
        from xmris.processing.utils import to_real_imag as _to_real_imag

        return _to_real_imag(self._obj, dim=dim, coords=coords)

    def to_complex(
        self, dim: str = DIMS.component, coords: tuple[str, str] = ("real", "imag")
    ) -> xr.DataArray:
        """Reconstruct a real-valued split array back into a standard complex array."""
        from xmris.processing.utils import to_complex as _to_complex

        return _to_complex(self._obj, dim=dim, coords=coords)
