from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import numpy as np
    from numpy.typing import ArrayLike, NDArray
else:
    ArrayLike = Any  # type: ignore[misc,assignment]
    NDArray = Any  # type: ignore[misc,assignment]


def _as_1d_float_array(values: ArrayLike, *, name: str) -> NDArray[np.float64]:
    import numpy as np

    array = np.asarray(values, dtype=np.float64)
    if array.ndim != 1:
        raise ValueError(f"{name} must be 1D, got shape={array.shape}")
    return array


@dataclass(frozen=True, slots=True)
class B1AFI:
    """Actual Flip Angle Imaging (AFI) for B1+ mapping (qMRLab: b1_afi).

    Protocol:
        - nom_fa_deg: nominal excitation flip angle [deg]
        - tr1_ms, tr2_ms: repetition times (any consistent units; only ratio matters)

    Input signal ordering:
        - signal = [AFIData1, AFIData2]
          where AFIData1 corresponds to TR1 and AFIData2 corresponds to TR2.
    """

    nom_fa_deg: float
    tr1_ms: float
    tr2_ms: float

    def __post_init__(self) -> None:
        if self.nom_fa_deg <= 0:
            raise ValueError("nom_fa_deg must be > 0")
        if self.tr1_ms <= 0 or self.tr2_ms <= 0:
            raise ValueError("tr1_ms and tr2_ms must be > 0")

    def fit(self, signal: ArrayLike) -> dict[str, float]:
        """Fit B1 from AFI signals (alias of ``fit_raw``)."""
        return self.fit_raw(signal)

    def fit_raw(self, signal: ArrayLike) -> dict[str, float]:
        """Fit B1 from AFI signals [S(TR1), S(TR2)].

        Parameters
        ----------
        signal : array-like
            Signal array ``[S(TR1), S(TR2)]``.

        Returns
        -------
        dict
            ``b1_raw`` and ``spurious`` (1.0 if b1_raw < 0.5).
        """
        import numpy as np

        y = _as_1d_float_array(signal, name="signal")
        if y.shape != (2,):
            raise ValueError("signal must be shape (2,) as [AFIData1, AFIData2]")

        s1 = float(y[0])
        s2 = float(y[1])
        if abs(s1) < 1e-12:
            return {"b1_raw": float("nan"), "spurious": 1.0}

        n = float(self.tr2_ms) / float(self.tr1_ms)
        r = abs(s2 / s1)

        # qMRLab behavior: r>1 treated as noise
        cos_arg = (r * n - 1.0) / (n - r) if r != n else 1.0
        if r > 1.0:
            cos_arg = 1.0

        cos_arg = float(np.clip(cos_arg, -1.0, 1.0))
        afi_rad = float(np.arccos(cos_arg))
        afi_deg = afi_rad * 180.0 / float(np.pi)
        b1_raw = afi_deg / float(self.nom_fa_deg)

        spurious = 1.0 if (not np.isfinite(b1_raw) or b1_raw < 0.5) else 0.0
        return {"b1_raw": float(b1_raw), "spurious": float(spurious)}

    def fit_image(
        self,
        signal: ArrayLike,
        *,
        mask: ArrayLike | str | None = None,
        n_jobs: int = 1,
        verbose: bool = False,
    ) -> dict[str, Any]:
        """Vectorized AFI B1 estimation on an image/volume.

        Parameters
        ----------
        signal : array-like
            Input array with last dim 2 as ``[AFIData1, AFIData2]``.
        mask : array-like, optional
            Spatial mask. If "otsu", Otsu thresholding is applied.
        n_jobs : int, default=1
            Number of parallel jobs. -1 uses all CPUs.
            Note: This model is fully vectorized and does not use n_jobs.
        verbose : bool, default=False
            If True, log info.

        Returns
        -------
        dict
            Maps for ``b1_raw`` and ``spurious``.
        """
        import logging

        import numpy as np

        from qmrpy._mask import resolve_mask

        logger = logging.getLogger("qmrpy")

        arr = np.asarray(signal, dtype=np.float64)
        if arr.ndim == 1:
            if mask is not None:
                raise ValueError("mask must be None for 1D data")
            return self.fit(arr)
        if arr.shape[-1] != 2:
            raise ValueError("data must have last dim=2 as [AFIData1, AFIData2]")

        spatial_shape = arr.shape[:-1]
        s1 = arr[..., 0]
        s2 = arr[..., 1]

        resolved_mask = resolve_mask(mask, arr)
        if resolved_mask is None:
            m = np.ones(spatial_shape, dtype=bool)
        else:
            if resolved_mask.shape != spatial_shape:
                raise ValueError(f"mask shape {resolved_mask.shape} must match spatial shape {spatial_shape}")
            m = resolved_mask

        b1_raw = np.full(spatial_shape, np.nan, dtype=np.float64)
        spurious = np.ones(spatial_shape, dtype=np.float64)

        valid = m & np.isfinite(s1) & np.isfinite(s2) & (np.abs(s1) >= 1e-12)
        n_voxels = int(np.sum(valid))
        n = float(self.tr2_ms) / float(self.tr1_ms)

        if verbose:
            logger.info("B1AFI: %d voxels (vectorized), shape=%s", n_voxels, spatial_shape)

        r = np.empty_like(s1, dtype=np.float64)
        r[valid] = np.abs(s2[valid] / s1[valid])

        cos_arg = np.full_like(r, 1.0, dtype=np.float64)
        denom = n - r
        mask_ok = valid & (denom != 0)
        cos_arg[mask_ok] = (r[mask_ok] * n - 1.0) / denom[mask_ok]
        cos_arg = np.clip(cos_arg, -1.0, 1.0)
        cos_arg = np.where(r > 1.0, 1.0, cos_arg)

        afi_deg = np.arccos(cos_arg) * 180.0 / float(np.pi)
        b1 = afi_deg / float(self.nom_fa_deg)

        b1_raw[valid] = b1[valid]
        spurious[valid] = np.where(np.isfinite(b1[valid]) & (b1[valid] >= 0.5), 0.0, 1.0)

        if verbose:
            logger.info("B1AFI complete: %d voxels processed", n_voxels)

        return {"b1_raw": b1_raw, "spurious": spurious}
