"""Xarray accessor for fUSI-specific operations."""

import xarray as xr

from confusius._utils import _compute_origin, _compute_spacing
from confusius.xarray.extract import FUSIExtractAccessor
from confusius.xarray.io import FUSIIOAccessor
from confusius.xarray.iq import FUSIIQAccessor
from confusius.xarray.plotting import FUSIPlotAccessor
from confusius.xarray.registration import FUSIRegistrationAccessor
from confusius.xarray.scale import FUSIScaleAccessor

__all__ = ["FUSIAccessor"]


@xr.register_dataarray_accessor("fusi")
class FUSIAccessor:
    """Xarray accessor for fUSI-specific operations.

    Provides convenient methods for functional ultrasound imaging data analysis.

    Parameters
    ----------
    xarray_obj : xarray.DataArray
        The `DataArray` to wrap.

    Examples
    --------
    >>> import xarray as xr
    >>> import numpy as np
    >>> from confusius import xarray as cxr  # Registers the accessor
    >>> data = xr.DataArray([1, 10, 100, 1000])
    >>> data.fusi.scale.db(factor=20)
    <xarray.DataArray (dim_0: 4)>
    array([-60., -40., -20.,   0.])
    """

    def __init__(self, xarray_obj: xr.DataArray) -> None:
        self._obj = xarray_obj

    @property
    def scale(self) -> FUSIScaleAccessor:
        """Access scaling operations.

        Returns
        -------
        FUSIScaleAccessor
            Accessor for scaling transformations.

        Examples
        --------
        >>> data = xr.DataArray([1, 10, 100, 1000])
        >>> data.fusi.scale.db(factor=10)
        <xarray.DataArray (dim_0: 4)>
        array([-30., -20., -10.,   0.])
        """
        return FUSIScaleAccessor(self._obj)

    @property
    def plot(self) -> FUSIPlotAccessor:
        """Access plotting operations.

        Returns
        -------
        FUSIPlotAccessor
            Accessor for plotting methods.

        Examples
        --------
        >>> import xarray as xr
        >>> data = xr.open_zarr("output.zarr")["iq"]
        >>> viewer, layer = data.fusi.plot.napari()
        """
        return FUSIPlotAccessor(self._obj)

    @property
    def register(self) -> FUSIRegistrationAccessor:
        """Access registration operations.

        Returns
        -------
        FUSIRegistrationAccessor
            Accessor for registration methods.

        Examples
        --------
        >>> import xarray as xr
        >>> data = xr.open_zarr("output.zarr")["pwd"]
        >>> registered = data.fusi.register.volumewise(reference_time=0)
        """
        return FUSIRegistrationAccessor(self._obj)

    @property
    def io(self) -> FUSIIOAccessor:
        """Access IO operations.

        Returns
        -------
        FUSIIOAccessor
            Accessor for IO methods (to_nii, etc.).

        Examples
        --------
        >>> import xarray as xr
        >>> data = xr.open_zarr("output.zarr")["pwd"]
        >>> data.fusi.io.to_nii("recording.nii.gz")
        """
        return FUSIIOAccessor(self._obj)

    @property
    def iq(self) -> FUSIIQAccessor:
        """Access IQ processing operations.

        Returns
        -------
        FUSIIQAccessor
            Accessor for IQ processing methods.

        Examples
        --------
        >>> import xarray as xr
        >>> ds = xr.open_zarr("output.zarr")
        >>> iq = ds["iq"]
        >>> pwd = iq.fusi.iq.process_to_power_doppler()
        >>> velocity = iq.fusi.iq.process_to_axial_velocity()
        """
        return FUSIIQAccessor(self._obj)

    @property
    def extract(self) -> FUSIExtractAccessor:
        """Access signal extraction operations.

        Returns
        -------
        FUSIExtractAccessor
            Accessor for extracting signals from fUSI data and reconstructing fUSI data
            from processed signals.

        Examples
        --------
        >>> import xarray as xr
        >>> data = xr.open_zarr("output.zarr")["pwd"]
        >>> mask = xr.open_zarr("brain_mask.zarr")["mask"]
        >>> signals = data.fusi.extract.with_mask(mask)
        >>> # ... process signals ...
        >>> restored = signals.fusi.extract.unmask(mask)
        """
        return FUSIExtractAccessor(self._obj)

    @property
    def spacing(self) -> dict[str, float | None]:
        """Coordinate spacing for all dimensions.

        Spacing is computed as the median of consecutive coordinate differences.
        A coordinate is considered uniform if its relative interval range
        `(max_diff - min_diff) / median_diff` is below 1%.

        Returns
        -------
        dict[str, float | None]
            Spacing per dimension, in DataArray dimension order. Returns `None` for
            dimensions with non-uniform or undefined spacing, with a warning.

        Examples
        --------
        >>> import xarray as xr
        >>> import numpy as np
        >>> import confusius  # noqa: F401
        >>> data = xr.DataArray(
        ...     np.zeros((10, 20)),
        ...     dims=["y", "x"],
        ...     coords={"y": np.arange(10) * 0.2, "x": np.arange(20) * 0.1},
        ... )
        >>> data.fusi.spacing
        {'y': 0.2, 'x': 0.1}
        """
        return _compute_spacing(self._obj)

    @property
    def origin(self) -> dict[str, float]:
        """Physical origin (first coordinate value) for all dimensions.

        For each dimension, returns the first coordinate value. If a coordinate is
        missing, warns and falls back to `0.0`.

        Returns
        -------
        dict[str, float]
            Origin per dimension, in DataArray dimension order.

        Examples
        --------
        >>> import xarray as xr
        >>> import numpy as np
        >>> import confusius  # noqa: F401
        >>> data = xr.DataArray(
        ...     np.zeros((10, 20)),
        ...     dims=["y", "x"],
        ...     coords={"y": np.arange(10) * 0.2, "x": np.arange(20) * 0.1},
        ... )
        >>> data.fusi.origin
        {'y': 0.0, 'x': 0.0}
        """
        return _compute_origin(self._obj)
