"""
pyamica._mne - MNE-Python integration for PyTorch AMICA.

Provides AmicaICA, a wrapper that fits AMICA on MNE Raw / Epochs data and
integrates with the full MNE artifact-rejection workflow.

Single-model usage (M=1)
------------------------
    ica = AmicaICA(device='cuda')
    ica.fit(raw, picks='eeg')

    ica.plot_components(inst=raw)
    ica.plot_sources(raw, block=True)   # click to mark bad components
    ica.apply(raw)                      # removes marked components in-place

Multi-model usage (M>1)
-----------------------
    ica = AmicaICA(n_models=3, device='cuda', chunk_t=32768)
    ica.fit(raw, picks='eeg')

    ica.plot_model_posteriors()         # inspect which model dominates when

    for m in range(3):
        ica.plot_sources(raw, model_idx=m, block=True)  # mark bad ICs per model

    ica.apply(raw)   # applies per-model exclusions using hard posteriors assignment
"""
from __future__ import annotations

import numpy as np
import torch

from pyamica._core import AMICA


# ─────────────────────────────────────────────────────────────────────────────
# Internal helpers
# ─────────────────────────────────────────────────────────────────────────────

def _resolve_picks(info, picks):
    """Return an int array of channel indices for *picks*."""
    import mne
    if isinstance(picks, str):
        _type_map = {
            'eeg':  dict(eeg=True,      meg=False),
            'meg':  dict(meg=True,      eeg=False),
            'grad': dict(meg='grad',    eeg=False),
            'mag':  dict(meg='mag',     eeg=False),
        }
        kw = _type_map.get(picks, {picks: True})
        return mne.pick_types(info, exclude='bads', **kw)
    elif isinstance(picks, (list, np.ndarray)):
        first = picks[0] if len(picks) else None
        if isinstance(first, str):
            return mne.pick_channels(info['ch_names'], picks,
                                     exclude=info['bads'])
        return np.asarray(picks, dtype=int)
    raise TypeError(f"picks must be str or list, got {type(picks).__name__}")


# ─────────────────────────────────────────────────────────────────────────────
# Standalone scoring helper (works with any fitted mne.preprocessing.ICA)
# ─────────────────────────────────────────────────────────────────────────────

def score_dipolarity(mne_ica, inst) -> 'np.ndarray':
    """
    Score each IC by how well its topomap resembles a single dipole.

    A dipolar scalp field varies roughly linearly across the electrode
    array. This function fits a linear regression of each IC topomap onto
    the x/y channel coordinates and returns the coefficient of
    determination (R squared). Values close to 1 indicate a dipolar
    pattern; values near 0 indicate a spatially diffuse or noisy map.

    This is a fast proxy for dipolarity: a linear fit on x/y channel
    positions captures the dominant spatial gradient of a single dipole
    without requiring a full head model.

    Works with any fitted ``mne.preprocessing.ICA`` object, regardless of
    the decomposition method (AMICA, FastICA, extended Infomax, etc.).

    Parameters
    ----------
    mne_ica : mne.preprocessing.ICA
        A fitted MNE ICA object (any method).
    inst : mne.io.BaseRaw or mne.BaseEpochs or mne.Evoked
        Used only for channel position information. Channels with missing
        or invalid positions (zero norm or NaN) are silently excluded from
        the regression. An error is only raised when no valid positions
        remain at all.

    Returns
    -------
    scores : np.ndarray, shape (n_components,)
        R squared of the linear fit for each IC topomap.
        Higher values indicate more dipolar patterns.

    Raises
    ------
    ValueError
        If no EEG channels with valid positions are found. This happens
        when no montage has been applied, since MNE initialises all
        channel locations to zero by default. Apply a montage first,
        e.g. ``inst.set_montage('standard_1020')``.
    """
    import numpy as np
    import mne

    topomaps = mne_ica.get_components()          # (n_ch, n_comp)

    picks = mne.pick_types(inst.info, eeg=True, exclude='bads')
    if len(picks) == 0:
        raise ValueError(
            "No EEG channels found. score_dipolarity() requires EEG "
            "channels with positions.")

    pos3d = np.array([inst.info['chs'][i]['loc'][:3] for i in picks])
    valid = np.isfinite(pos3d).all(axis=1) & (np.linalg.norm(pos3d, axis=1) > 1e-10)
    if not valid.any():
        raise ValueError(
            "No valid EEG channel positions found. Apply a montage before "
            "calling score_dipolarity(), e.g. "
            "inst.set_montage('standard_1020').")

    pos3d    = pos3d[valid]
    topomaps = topomaps[valid]

    norms = np.linalg.norm(pos3d, axis=1, keepdims=True)
    xy    = pos3d[:, :2] / np.where(norms > 0, norms, 1.0)
    X_reg = np.column_stack([xy, np.ones(len(pos3d))])   # (n_ch, 3)

    scores = np.empty(topomaps.shape[1])
    for i in range(topomaps.shape[1]):
        topo          = topomaps[:, i]
        beta, _, _, _ = np.linalg.lstsq(X_reg, topo, rcond=None)
        resid         = topo - X_reg @ beta
        ss_res        = float(np.dot(resid, resid))
        ss_tot        = float(np.dot(topo - topo.mean(), topo - topo.mean()))
        scores[i]     = 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0

    return scores


# ─────────────────────────────────────────────────────────────────────────────
# Mutual information helpers
# ─────────────────────────────────────────────────────────────────────────────

def _mi_matrix(data: 'np.ndarray', n_bins: int = 30) -> 'np.ndarray':
    """Pairwise mutual information matrix for a (n, T) data array (nats)."""
    import numpy as np
    n  = data.shape[0]
    mi = np.zeros((n, n))
    for i in range(n):
        for j in range(i + 1, n):
            c_xy = np.histogram2d(data[i], data[j], bins=n_bins)[0]
            p_xy = c_xy / c_xy.sum()
            p_x  = p_xy.sum(axis=1, keepdims=True)
            p_y  = p_xy.sum(axis=0, keepdims=True)
            mask = p_xy > 0
            val  = float(np.sum(p_xy[mask] * np.log(p_xy[mask] / (p_x * p_y)[mask])))
            mi[i, j] = mi[j, i] = max(0.0, val)
    return mi


def score_mutual_information(mne_ica, inst, n_bins: int = 30) -> 'np.ndarray':
    """
    Compute the pairwise mutual information matrix between IC time series.

    Mutual information (MI) measures statistical dependence between two
    signals: zero means complete independence, larger values indicate
    stronger dependence. ICA aims to minimise MI between components, so
    a well-fitted decomposition should produce a matrix that is close to
    zero everywhere off the diagonal.

    MI is estimated from a joint histogram with *n_bins* per dimension.
    The estimator is biased towards zero for small T; use at least
    T = 1000 samples for reliable values.

    Works with any fitted ``mne.preprocessing.ICA`` object, regardless of
    the decomposition method (AMICA, Infomax, extended Infomax, etc.).

    Parameters
    ----------
    mne_ica : mne.preprocessing.ICA
        A fitted MNE ICA object (any method).
    inst : mne.io.BaseRaw or mne.BaseEpochs
        Data to extract IC time series from.
    n_bins : int
        Number of histogram bins per dimension. Default 30.

    Returns
    -------
    mi : np.ndarray, shape (n_components, n_components)
        Symmetric matrix of pairwise MI values in nats. Diagonal is zero.
    """
    sources = mne_ica.get_sources(inst).get_data()   # (n_comp, T)
    return _mi_matrix(sources, n_bins=n_bins)


# ─────────────────────────────────────────────────────────────────────────────
# Main class
# ─────────────────────────────────────────────────────────────────────────────

class AmicaICA:
    """
    AMICA ICA fitted to MNE Raw or Epochs data.

    Parameters
    ----------
    n_components : int or None
        Number of ICA components. Default None (equal to the number of
        selected channels).
    n_models : int
        Number of AMICA mixture models (M). 1 = standard ICA. M > 1 learns
        separate unmixing matrices and source densities per mixture model. Default 1.
    device : str
        PyTorch device string: ``'cpu'``, ``'cuda'``, or ``'mps'``.
        Default ``'cpu'``.
    **amica_kwargs
        Passed verbatim to :class:`~pyamica.AMICA`, e.g. ``max_iter``,
        ``lrate``, ``do_newton``, ``compile``, ``time_iters``, ``chunk_t``,
        ``do_reject``, ``reject_sigma``, ``num_reject``, ``reject_start``,
        ``reject_int``.

    Notes
    -----
    Data are scaled from Volts (MNE convention) to microvolts before fitting.
    This improves numerical conditioning of the PCA/sphering step and matches
    the scale on which AMICA's default hyperparameters were designed.
    The MNE ICA object returned by :meth:`to_mne_ica` operates in Volts so it
    integrates transparently with the rest of MNE.
    """

    _SCALE: float = 1e6   # V → µV

    def __init__(self,
                 n_components: int | None = None,
                 n_models: int = 1,
                 device: str = 'cpu',
                 **amica_kwargs):
        self.n_components = n_components
        self.n_models     = n_models
        self.device       = device
        self.amica_kwargs = amica_kwargs

        # populated by fit()
        self._model:      AMICA | None      = None
        self._picks_idx:  np.ndarray | None = None
        self._ch_names:   list[str] | None  = None
        self._inst_info                     = None   # mne.Info for picked channels
        self._n_samples:  int               = 0
        self._fit_type:   str               = 'raw'  # 'raw' | 'epc'
        self._times:      np.ndarray | None = None   # (T,) timestamps of good samples (s)
        self._mne_icas:   dict              = {}     # cache: model_idx → MNE ICA

    # ─────────────────────────────────────────────────────────────────────────

    def fit(self, inst, picks: str | list = 'eeg') -> 'AmicaICA':
        """
        Fit AMICA on the good data in *inst*.

        Bad channels (``inst.info['bads']``) and bad time segments
        (BAD annotations in Raw; rejected Epochs) are automatically excluded
        via MNE's built-in data extraction.

        Parameters
        ----------
        inst : mne.io.BaseRaw or mne.BaseEpochs
        picks : str or list of str
            Channel type string ('eeg', 'meg', 'grad', 'mag') or explicit
            list of channel names.

        Returns
        -------
        self
        """
        import mne

        # ── channel selection ──────────────────────────────────────────────
        picks_idx = _resolve_picks(inst.info, picks)
        if len(picks_idx) == 0:
            raise ValueError(f"No channels selected by picks={picks!r}.")
        self._picks_idx = picks_idx
        self._ch_names  = [inst.info['ch_names'][i] for i in picks_idx]
        self._inst_info = mne.pick_info(inst.info, picks_idx)

        # ── extract good data (MNE handles bad segments / epochs) ──────────
        if isinstance(inst, mne.io.BaseRaw):
            # reject_by_annotation='omit' concatenates only the good segments
            data, times = inst.get_data(picks=picks_idx,
                                        reject_by_annotation='omit',
                                        return_times=True)          # (n_ch, T), (T,)
            self._times    = times
            self._fit_type = 'raw'
        elif isinstance(inst, mne.BaseEpochs):
            # get_data() returns only non-rejected epochs
            ep = inst.get_data(picks=picks_idx)                    # (n_ep, n_ch, n_t)
            n_ep, n_ch, n_t = ep.shape
            data = ep.transpose(0, 2, 1).reshape(n_ep * n_t, n_ch).T  # (n_ch, T)
            # unroll epoch-relative times across all epochs
            self._times    = np.tile(inst.times, n_ep)
            self._fit_type = 'epc'
        else:
            raise TypeError(
                f"inst must be mne.io.BaseRaw or mne.BaseEpochs, "
                f"got {type(inst).__name__}"
            )

        # ── V → µV (better numerical conditioning for covariance / sphere) ─
        data = data * self._SCALE                                   # (n_ch, T) µV

        # ── (n_ch, T) → (T, n_ch)  (AMICA expects time-first) ────────────
        X = torch.from_numpy(data.T.astype('float64'))             # (T, n_ch)
        self._n_samples = X.shape[0]

        # ── fit AMICA ──────────────────────────────────────────────────────
        self._mne_icas = {}   # invalidate any previously cached MNE ICA objects
        n_comp = self.n_components or data.shape[0]
        self._model = AMICA(
            n_components = n_comp,
            n_models     = self.n_models,
            device       = self.device,
            **self.amica_kwargs,
        )
        self._model.fit(X)
        return self

    # ─────────────────────────────────────────────────────────────────────────

    def plot_model_posteriors(self, ax=None, figsize=None):
        """
        Plot per-time-point model posteriors p(m|t) against recording time.

        Each line shows the probability that model m is active at each sample.
        For well-separated conditions the lines should approach 0/1 in distinct
        time segments, making it easy to verify that AMICA's learned models
        align with the known experimental conditions.

        Parameters
        ----------
        ax : matplotlib.axes.Axes or None
            Axes to plot into.  If None, a new figure is created.
        figsize : tuple or None
            Figure size passed to ``plt.figure()``.  Ignored if *ax* is given.

        Returns
        -------
        matplotlib.axes.Axes
        """
        import matplotlib.pyplot as plt

        if self._model is None or self._model.posteriors_ is None:
            raise RuntimeError("Call fit() before plot_model_posteriors().")
        if self._times is None:
            raise RuntimeError("No time information available. fit() may not have stored it.")

        assert self._model.gm_ is not None
        posteriors = self._model.posteriors_.cpu().numpy()   # (M, T)
        gm         = self._model.gm_.cpu().numpy()           # (M,)
        times      = self._times                             # (T,)
        M          = posteriors.shape[0]

        if ax is None:
            fig, ax = plt.subplots(figsize=figsize or (12, 3))

        colors = plt.cm.tab10.colors
        for m in range(M):
            ax.plot(times, posteriors[m],
                    color=colors[m % len(colors)],
                    lw=0.5, alpha=0.8,
                    label=f'Model {m}  (gm={gm[m]:.3f})')

        ax.set_xlabel('Time (s)')
        ax.set_ylabel('p(model | t)')
        ax.set_title('AMICA model posteriors')
        ax.set_ylim(-0.05, 1.05)
        ax.legend(loc='upper right')
        ax.figure.tight_layout()
        return ax

    # ─────────────────────────────────────────────────────────────────────────

    def plot_model_dominance(self, smooth_s: float = 0.0,
                             figsize=None, ax=None):
        """
        Stacked area chart of model posteriors, easier to read than
        overlapping lines when models switch rapidly.

        Because posteriors sum to 1, each model's share is shown as a filled
        band.  Rapid within-second switches appear as mixed colours rather
        than unreadable flickering.

        An optional Gaussian smooth (``smooth_s`` seconds) can be applied
        before plotting to suppress fast transients and reveal the broader
        structure of model dominance.

        Parameters
        ----------
        smooth_s : float
            Standard deviation of the Gaussian kernel (seconds).
            ``0`` (default) means no smoothing.
        ax : matplotlib.axes.Axes or None
            Axes to plot into.  If None, a new figure is created.
        figsize : tuple or None
            Figure size.  Ignored when *ax* is given.

        Returns
        -------
        matplotlib.axes.Axes
        """
        import matplotlib.pyplot as plt
        from scipy.ndimage import gaussian_filter1d

        if self._model is None or self._model.posteriors_ is None:
            raise RuntimeError("Call fit() before plot_model_dominance().")
        if self._times is None:
            raise RuntimeError("No time information available.")

        assert self._model.gm_ is not None
        posteriors = self._model.posteriors_.cpu().numpy().copy()   # (M, T)
        gm         = self._model.gm_.cpu().numpy()
        times      = self._times
        M          = posteriors.shape[0]

        if smooth_s > 0:
            sfreq  = 1.0 / (times[1] - times[0]) if len(times) > 1 else 1.0
            sigma  = smooth_s * sfreq              # samples
            posteriors = gaussian_filter1d(posteriors, sigma=sigma, axis=1)
            # re-normalise after smoothing (posteriors should already sum to ~1)
            posteriors /= posteriors.sum(axis=0, keepdims=True).clip(min=1e-30)

        if ax is None:
            fig, ax = plt.subplots(figsize=figsize or (12, 3))

        colors = [plt.cm.tab10.colors[m % 10] for m in range(M)]
        labels = [f'Model {m}  (gm={gm[m]:.3f})' for m in range(M)]

        ax.stackplot(times, posteriors, labels=labels, colors=colors, alpha=0.85)

        smooth_str = f', smoothed {smooth_s}s' if smooth_s > 0 else ''
        ax.set_xlabel('Time (s)')
        ax.set_ylabel('p(model | t)')
        ax.set_title(f'AMICA model dominance{smooth_str}')
        ax.set_ylim(0, 1)
        ax.legend(loc='upper right')
        ax.figure.tight_layout()
        return ax

    # ─────────────────────────────────────────────────────────────────────────

    def get_mne_ica(self, model_idx: int = 0):
        """
        Return the cached ``mne.preprocessing.ICA`` for *model_idx*.

        The object is created on first access and cached; subsequent calls
        return the same object so that interactive exclusion selections
        (``ica.exclude``) made in ``plot_sources`` / ``plot_components``
        are preserved and picked up by ``apply()``.

        Parameters
        ----------
        model_idx : int
            0-indexed model.  Model 0 is always the most probable model after fit().
        """
        if self._model is None:
            raise RuntimeError("Call fit() before get_mne_ica().")
        if model_idx not in self._mne_icas:
            self._mne_icas[model_idx] = self.to_mne_ica(model_idx)
        return self._mne_icas[model_idx]

    # ─────────────────────────────────────────────────────────────────────────

    def plot_sources(self, inst, model_idx: int = 0, **kwargs):
        """
        Plot ICA source time-courses for *model_idx* (delegates to MNE).

        Clicking on a source in the interactive window marks it for removal;
        selections are stored on the cached MNE ICA object and read by
        ``apply()``.

        Parameters
        ----------
        inst : mne.io.BaseRaw or mne.BaseEpochs
        model_idx : int
            Which AMICA model to inspect.  Default 0.
        **kwargs
            Forwarded to ``mne.preprocessing.ICA.plot_sources()``.
        """
        return self.get_mne_ica(model_idx).plot_sources(inst, **kwargs)

    # ─────────────────────────────────────────────────────────────────────────

    def plot_components(self, inst=None, model_idx: int = 0, **kwargs):
        """
        Plot ICA component topographies for *model_idx* (delegates to MNE).

        Parameters
        ----------
        inst : mne.io.BaseRaw or mne.BaseEpochs or None
        model_idx : int
            Which AMICA model to inspect.  Default 0.
        **kwargs
            Forwarded to ``mne.preprocessing.ICA.plot_components()``.
        """
        return self.get_mne_ica(model_idx).plot_components(inst=inst, **kwargs)

    # ─────────────────────────────────────────────────────────────────────────

    def plot_properties(self, inst, picks=None, model_idx: int = 0, **kwargs):
        """
        Plot detailed component properties for *model_idx* (delegates to MNE).

        Parameters
        ----------
        inst : mne.io.BaseRaw or mne.BaseEpochs
        picks : int or list of int or None
            Component indices to plot.  None = all.
        model_idx : int
            Which AMICA model to inspect.  Default 0.
        **kwargs
            Forwarded to ``mne.preprocessing.ICA.plot_properties()``.
        """
        return self.get_mne_ica(model_idx).plot_properties(inst, picks=picks, **kwargs)

    # ─────────────────────────────────────────────────────────────────────────

    def plot_overlay(self, inst, model_idx: int = 0, **kwargs):
        """
        Plot original vs. ICA-cleaned signal for *model_idx* (delegates to MNE).

        Uses the excluded components stored on the cached MNE ICA object, so
        call this after marking components via ``plot_sources()`` to preview
        what will be removed.

        Parameters
        ----------
        inst : mne.io.BaseRaw or mne.Evoked
        model_idx : int
            Which AMICA model to inspect.  Default 0.
        **kwargs
            Forwarded to ``mne.preprocessing.ICA.plot_overlay()``.
        """
        return self.get_mne_ica(model_idx).plot_overlay(inst, **kwargs)

    # ─────────────────────────────────────────────────────────────────────────

    def plot_scores(self, scores, model_idx: int = 0, **kwargs):
        """
        Plot component scores (e.g. from ``find_bads_eog``) for *model_idx*.

        Typical usage::

            eog_idx, scores = ica.find_bads_eog(raw, model_idx=m, ch_name='HEOG')
            ica.plot_scores(scores, model_idx=m)

        Parameters
        ----------
        scores : array-like
            Scores returned by ``find_bads_eog()`` or ``find_bads_ecg()``.
        model_idx : int
            Which AMICA model the scores belong to.  Default 0.
        **kwargs
            Forwarded to ``mne.preprocessing.ICA.plot_scores()``.
        """
        return self.get_mne_ica(model_idx).plot_scores(scores, **kwargs)

    # ─────────────────────────────────────────────────────────────────────────

    def find_bads_eog(self, inst, model_idx: int = 0, **kwargs):
        """
        Identify EOG-correlated components for *model_idx* (delegates to MNE).

        Parameters
        ----------
        inst : mne.io.BaseRaw or mne.BaseEpochs
        model_idx : int
            Which AMICA model to inspect.  Default 0.
        **kwargs
            Forwarded to ``mne.preprocessing.ICA.find_bads_eog()``
            (e.g. ``ch_name``, ``threshold``, ``measure``).

        Returns
        -------
        eog_indices : list of int
        scores : ndarray
        """
        return self.get_mne_ica(model_idx).find_bads_eog(inst, **kwargs)

    # ─────────────────────────────────────────────────────────────────────────

    def find_bads_ecg(self, inst, model_idx: int = 0, **kwargs):
        """
        Identify ECG-correlated components for *model_idx* (delegates to MNE).

        Parameters
        ----------
        inst : mne.io.BaseRaw or mne.BaseEpochs
        model_idx : int
            Which AMICA model to inspect.  Default 0.
        **kwargs
            Forwarded to ``mne.preprocessing.ICA.find_bads_ecg()``
            (e.g. ``ch_name``, ``threshold``, ``method``).

        Returns
        -------
        ecg_indices : list of int
        scores : ndarray
        """
        return self.get_mne_ica(model_idx).find_bads_ecg(inst, **kwargs)

    # ─────────────────────────────────────────────────────────────────────────

    def review(self, inst, model_idx: int = 0,
               eog_ch=None, ecg_ch=None,
               eog_kwargs: dict | None = None,
               ecg_kwargs: dict | None = None) -> list[int]:
        """
        Semi-automatic component review for one model.

        Runs automated bad-component detection, prints a terminal summary,
        pre-selects flagged components, then opens ``plot_sources()`` so you
        can confirm or adjust the selection interactively before closing the
        window.

        Parameters
        ----------
        inst : mne.io.BaseRaw or mne.BaseEpochs
        model_idx : int
            Which AMICA model to review.  Default 0.
        eog_ch : str or list of str or None
            EOG channel name(s) to pass to ``find_bads_eog()``.
            Pass a list to run detection separately for each channel
            (e.g. ``['HEOG', 'VEOG']``).
        ecg_ch : str or None
            ECG channel name for ``find_bads_ecg()``.
        eog_kwargs : dict or None
            Extra keyword arguments for ``find_bads_eog()``.
        ecg_kwargs : dict or None
            Extra keyword arguments for ``find_bads_ecg()``.

        Returns
        -------
        list of int
            Final component exclusion list for this model (same object as
            ``get_mne_ica(model_idx).exclude``).

        Example
        -------
        ::

            for m in range(ica.n_models):
                ica.review(raw, model_idx=m,
                           eog_ch=['HEOG', 'VEOG'], ecg_ch='ECG')
            ica.apply(raw)
        """
        mne_ica = self.get_mne_ica(model_idx)
        # flagged: label → (indices, scores_array of shape (n_comp,))
        flagged: dict[str, tuple[list[int], np.ndarray]] = {}

        # ── automated detection ───────────────────────────────────────────────
        if eog_ch is not None:
            channels = [eog_ch] if isinstance(eog_ch, str) else list(eog_ch)
            for ch in channels:
                idx, scores = mne_ica.find_bads_eog(inst, ch_name=ch,
                                                     **(eog_kwargs or {}))
                flagged[f'EOG ({ch})'] = (idx, scores)

        if ecg_ch is not None:
            idx, scores = mne_ica.find_bads_ecg(inst, ch_name=ecg_ch,
                                                 **(ecg_kwargs or {}))
            flagged[f'ECG ({ecg_ch})'] = (idx, scores)

        # ── terminal summary ──────────────────────────────────────────────────
        print(f"\n{'='*60}")
        print(f"  Model {model_idx} - automated detection")
        print(f"{'='*60}")

        all_flagged: set[int] = set()
        if flagged:
            for label, (indices, scores) in flagged.items():
                all_flagged.update(indices)

                # scores for flagged components
                flagged_str = '  '.join(
                    f'IC{i:03d}={scores[i]:+.2f}' for i in sorted(indices)
                )

                # highest-scoring non-flagged component as a clean reference
                ranked = np.argsort(np.abs(scores))[::-1]
                clean  = next((i for i in ranked if i not in indices), None)
                ref_str = (f'  |  best clean: IC{clean:03d}={scores[clean]:+.2f}'
                           if clean is not None else '')

                print(f"  {label:<22s}: {flagged_str or '(none)'}{ref_str}")
        else:
            print("  (no automated detection requested)")

        mne_ica.exclude = sorted(all_flagged)
        print(f"\n  Pre-selected : {mne_ica.exclude}")
        print("  Close the sources window to confirm (add/remove by clicking).")

        # ── interactive review ────────────────────────────────────────────────
        self.plot_sources(inst, model_idx=model_idx, block=True)

        print(f"  Final exclusions for model {model_idx}: {mne_ica.exclude}")
        return list(mne_ica.exclude)

    # ─────────────────────────────────────────────────────────────────────────

    def score_dipolarity(self, inst, model_idx: int = 0) -> 'np.ndarray':
        """
        Score each IC by how well its topomap resembles a single dipole.

        A dipolar scalp field varies roughly linearly across the electrode
        array. This method fits a linear regression of each IC topomap onto
        the x/y channel coordinates and returns the coefficient of
        determination (R squared). Values close to 1 indicate a dipolar
        pattern; values near 0 indicate a spatially diffuse or noisy map.

        This is a fast proxy for dipolarity: a linear fit on x/y channel
        positions captures the dominant spatial gradient of a single dipole
        without requiring a full head model.

        Parameters
        ----------
        inst : mne.io.BaseRaw or mne.BaseEpochs or mne.Evoked
            Used only for channel position information. Channels with missing
            or invalid positions (zero norm or NaN, which MNE assigns when no
            montage is set) are silently excluded from the regression. An
            error is only raised when no valid positions remain at all.
        model_idx : int
            Which AMICA model to score. Default 0.

        Returns
        -------
        scores : np.ndarray, shape (n_components,)
            R squared of the linear fit for each IC topomap.
            Higher values indicate more dipolar patterns.

        Raises
        ------
        RuntimeError
            If fit() has not been called yet.
        ValueError
            If no EEG channels with valid positions are found. This happens
            when no montage has been applied, since MNE initialises all
            channel locations to zero by default. Apply a montage first,
            e.g. ``inst.set_montage('standard_1020')``.
        """
        if self._model is None:
            raise RuntimeError("Call fit() before score_dipolarity().")
        return score_dipolarity(self.get_mne_ica(model_idx), inst)

    # ─────────────────────────────────────────────────────────────────────────

    def score_mutual_information(
        self, inst, model_idx: int = 0, n_bins: int = 30
    ) -> 'np.ndarray':
        """
        Compute the pairwise mutual information matrix between IC time series.

        Mutual information (MI) measures statistical dependence between two
        signals: zero means complete independence, larger values indicate
        stronger dependence. ICA aims to minimise MI between components, so
        a well-fitted decomposition should produce a matrix that is close to
        zero everywhere off the diagonal.

        MI is estimated from a joint histogram with *n_bins* per dimension.
        The estimator is biased towards zero for small T; use at least
        T = 1000 samples for reliable values.

        Parameters
        ----------
        inst : mne.io.BaseRaw or mne.BaseEpochs
            Data to extract IC time series from.
        model_idx : int
            Which AMICA model to score. Default 0.
        n_bins : int
            Number of histogram bins per dimension. Default 30.

        Returns
        -------
        mi : np.ndarray, shape (n_components, n_components)
            Symmetric matrix of pairwise MI values in nats. Diagonal is zero.

        Raises
        ------
        RuntimeError
            If fit() has not been called yet.
        """
        if self._model is None:
            raise RuntimeError("Call fit() before score_mutual_information().")
        return score_mutual_information(self.get_mne_ica(model_idx), inst, n_bins)

    # ─────────────────────────────────────────────────────────────────────────

    def apply(self, inst) -> 'AmicaICA':
        """
        Apply ICA artifact removal in-place.

        Component exclusions are read from the cached MNE ICA objects. Set
        them interactively via ``plot_sources()`` / ``plot_components()``, or
        directly via ``get_mne_ica(m).exclude = [0, 2]``.

        For ``n_models=1`` this delegates directly to MNE's ``ICA.apply()``.
        For ``n_models>1`` a hard per-sample model assignment (argmax of
        posteriors) selects which model's unmixing and exclusion list is used
        at each time point.

        Parameters
        ----------
        inst : mne.io.BaseRaw or mne.BaseEpochs
            Modified **in-place**.

        Returns
        -------
        inst
        """
        if self._model is None:
            raise RuntimeError("Call fit() before apply().")

        M = self._model.n_models

        if M == 1:
            self.get_mne_ica(0).apply(inst)
        else:
            exclude_per_model = [list(self.get_mne_ica(m).exclude) for m in range(M)]
            self._apply_multi(inst, exclude_per_model)

        return inst

    # ─────────────────────────────────────────────────────────────────────────

    def _apply_multi(self, inst, exclude_per_model: list[list[int]]):
        """Hard-assignment multi-model apply (internal)."""
        import mne

        assert self._model is not None and self._model.posteriors_ is not None
        assert self._picks_idx is not None

        M          = self._model.n_models
        mne_icas   = [self.get_mne_ica(m) for m in range(M)]
        assignment = self._model.posteriors_.cpu().numpy().argmax(axis=0)
        picks_idx  = self._picks_idx

        if isinstance(inst, mne.io.BaseRaw):
            data = inst.get_data(picks=picks_idx)               # (n_ch, T_full)
            _, T_full = data.shape

            _, good_times = inst.get_data(picks=picks_idx[:1],
                                          reject_by_annotation='omit',
                                          return_times=True)
            good_idx = np.searchsorted(inst.times, good_times).clip(0, T_full - 1)

            full_assignment = np.zeros(T_full, dtype=int)
            full_assignment[good_idx] = assignment

            for m, mne_ica in enumerate(mne_icas):
                mask = full_assignment == m
                if not np.any(mask) or not exclude_per_model[m]:
                    continue
                mne_ica.exclude = exclude_per_model[m]
                seg_raw = mne.io.RawArray(
                    data[:, mask], mne.pick_info(inst.info, picks_idx), verbose=False
                )
                mne_ica.apply(seg_raw, verbose=False)
                data[:, mask] = seg_raw.get_data()

            inst._data[np.ix_(picks_idx, np.arange(T_full))] = data

        elif isinstance(inst, mne.BaseEpochs):
            ep_data = inst.get_data(picks=picks_idx)            # (n_ep, n_ch, n_t)
            n_ep, n_ch, n_t = ep_data.shape
            flat = ep_data.transpose(0, 2, 1).reshape(n_ep * n_t, n_ch).T

            for m, mne_ica in enumerate(mne_icas):
                mask = assignment == m
                if not np.any(mask) or not exclude_per_model[m]:
                    continue
                mne_ica.exclude = exclude_per_model[m]
                seg_raw = mne.io.RawArray(
                    flat[:, mask], mne.pick_info(inst.info, picks_idx), verbose=False
                )
                mne_ica.apply(seg_raw, verbose=False)
                flat[:, mask] = seg_raw.get_data()

            inst._data[:, picks_idx, :] = flat.T.reshape(n_ep, n_t, n_ch).transpose(0, 2, 1)

        else:
            raise TypeError(f"inst must be Raw or Epochs, got {type(inst).__name__}")

    # ─────────────────────────────────────────────────────────────────────────

    def save(self, path) -> None:
        """
        Save the fitted model and component selections to a ``.amica`` file.

        The file is a compressed NumPy archive containing all fitted tensors,
        recording metadata, and per-model exclusion lists.  Reload with
        ``AmicaICA.load(path)``.

        Parameters
        ----------
        path : str or Path
            Output path.  ``.amica.npz`` is appended if not already present.
        """
        import json
        import pickle
        from pathlib import Path as _Path

        if self._model is None:
            raise RuntimeError("Call fit() before save().")

        path = _Path(path)
        if not str(path).endswith('.amica.npz'):
            path = _Path(str(path) + '.amica.npz')

        m      = self._model
        arrays: dict[str, np.ndarray] = {}

        # ── AMICA tensors ─────────────────────────────────────────────────
        for attr in ['W_', 'A_', 'gm_', 'mu_', 'alpha_', 'beta_', 'rho_',
                     'sbeta_', 'mean_', 'sphere_', 'pca_vecs_', 'pca_vals_',
                     'posteriors_']:
            val = getattr(m, attr, None)
            if val is not None:
                arrays[f'model_{attr}'] = val.cpu().numpy()

        if self._times is not None:
            arrays['_times'] = self._times
        if self._picks_idx is not None:
            arrays['_picks_idx'] = self._picks_idx

        # ── metadata ──────────────────────────────────────────────────────
        meta = {
            'n_components': self.n_components,
            'n_models':     self.n_models,
            'device':       str(self.device),
            'ch_names':     self._ch_names,
            'n_samples':    self._n_samples,
            'fit_type':     self._fit_type,
            'sldet':        float(m.sldet_),
            'max_iter':     int(m.max_iter),
            'exclude_per_model': {
                str(mi): list(self._mne_icas[mi].exclude)
                for mi in self._mne_icas
            },
        }
        arrays['_meta'] = np.array(json.dumps(meta))

        # ── MNE info (pickled into a uint8 array) ─────────────────────────
        if self._inst_info is not None:
            arrays['_info_pkl'] = np.frombuffer(
                pickle.dumps(self._inst_info), dtype=np.uint8
            )

        np.savez_compressed(path, **arrays)   # path already ends in .npz, no suffix added
        print(f"Saved to {path}")

    # ─────────────────────────────────────────────────────────────────────────

    @classmethod
    def load(cls, path) -> 'AmicaICA':
        """
        Load a fitted AmicaICA from a ``.amica`` file saved with ``save()``.

        All fitted tensors and per-model exclusion lists are restored.  Call
        ``apply(inst)`` afterwards to clean the original recording.

        Parameters
        ----------
        path : str or Path
            Path to the ``.amica.npz`` file.

        Returns
        -------
        AmicaICA
        """
        import json
        import pickle
        from pathlib import Path as _Path

        path = _Path(path)
        data = np.load(path, allow_pickle=True)
        meta = json.loads(str(data['_meta']))

        # ── reconstruct AmicaICA ──────────────────────────────────────────
        obj              = cls.__new__(cls)
        obj.n_components = meta['n_components']
        obj.n_models     = meta['n_models']
        obj.device       = meta['device']
        obj.amica_kwargs = {}
        obj._ch_names    = meta['ch_names']
        obj._n_samples   = meta['n_samples']
        obj._fit_type    = meta['fit_type']
        obj._times       = data['_times']    if '_times'    in data else None
        obj._picks_idx   = data['_picks_idx'] if '_picks_idx' in data else None
        obj._mne_icas    = {}
        obj._inst_info   = (pickle.loads(data['_info_pkl'].tobytes())
                            if '_info_pkl' in data else None)

        # ── reconstruct minimal AMICA stub ────────────────────────────────
        # Only the attributes accessed by to_mne_ica() and apply() are needed.
        amica          = AMICA.__new__(AMICA)
        amica.n_models = meta['n_models']
        amica.max_iter = meta['max_iter']
        amica.sldet_   = float(meta['sldet'])

        device = torch.device('cpu')
        for attr in ['W_', 'A_', 'gm_', 'mu_', 'alpha_', 'beta_', 'rho_',
                     'sbeta_', 'mean_', 'sphere_', 'pca_vecs_', 'pca_vals_',
                     'posteriors_']:
            key = f'model_{attr}'
            setattr(amica, attr,
                    torch.from_numpy(data[key]).to(device) if key in data else None)

        obj._model = amica

        # ── restore exclusion lists ────────────────────────────────────────
        for mi_str, exclude in meta['exclude_per_model'].items():
            mi = int(mi_str)
            mne_ica         = obj.to_mne_ica(mi)
            mne_ica.exclude = list(exclude)
            obj._mne_icas[mi] = mne_ica

        n_comp = amica.W_.shape[-1] if amica.W_ is not None else '?'
        print(f"Loaded from {path}  "
              f"(n_models={meta['n_models']}, n_components={n_comp})")
        return obj

    # ─────────────────────────────────────────────────────────────────────────

    def to_mne_ica(self, model_idx: int = 0):
        """
        Build and return a fitted ``mne.preprocessing.ICA`` object.

        After fit(), AMICA models are sorted so that model 0 is always the most
        probable model (highest ``gm_``) and components within each model are
        ordered by decreasing variance explained.  For multi-model fits, call
        once per model:

            ica0 = amica_ica.to_mne_ica(model_idx=0)   # most probable model
            ica1 = amica_ica.to_mne_ica(model_idx=1)

        AMICA's ZCA sphere and W matrix are decomposed into MNE's internal
        PCA + ICA representation:

        * ``pca_components_``  = V.T  (orthonormal PCA eigenvectors)
        * ``unmixing_matrix_`` = W × V × diag(D⁻½)
        * ``mixing_matrix_``   = diag(D^½) × V.T × A

        where V, D come from the ZCA sphere (shared across all models) and
        A = inv(W) is AMICA's mixing matrix for the selected model.

        Parameters
        ----------
        model_idx : int
            Which AMICA model to export (0-indexed).  Default 0 (most probable
            model after fit()).

        Returns
        -------
        mne.preprocessing.ICA
            Fully populated, ready for ``plot_components()``,
            ``find_bads_eog()``, ``apply()``, ``save()`` / ``load()``.
        """
        import mne
        from mne.preprocessing import ICA as MNE_ICA

        if self._model is None or self._ch_names is None:
            raise RuntimeError("Call fit() before to_mne_ica().")

        m = self._model
        assert m.W_ is not None and m.pca_vecs_ is not None \
               and m.pca_vals_ is not None and m.mean_ is not None, \
               "AMICA model is missing fitted attributes. fit() may not have completed."

        if not (0 <= model_idx < m.n_models):
            raise ValueError(f"model_idx={model_idx} out of range for n_models={m.n_models}.")

        n_comp = int(m.W_.shape[1])    # W_ is (M, n_comp, n_comp)
        scale  = self._SCALE

        # ── Recover PCA eigenvectors / eigenvalues from stored attributes ──
        # pca_vecs_: (n_ch, n_comp) eigenvectors V  (set by _compute_sphere)
        # pca_vals_: (n_comp,) eigenvalues of cov in µV²  (descending)
        V        = m.pca_vecs_.cpu().numpy().astype('float64')     # (n_ch, n_comp)
        d_vals   = m.pca_vals_.cpu().numpy().astype('float64')     # µV²  descending

        d_invsqrt = 1.0 / np.sqrt(d_vals)                         # D^{-½}

        # ── AMICA matrices ─────────────────────────────────────────────────
        W = m.W_[model_idx].cpu().numpy().astype('float64')        # (n_comp, n_comp)
        A = m.A_[model_idx].cpu().numpy().astype('float64')        # (n_comp, n_comp) = inv(W)

        # ── MNE attribute mapping ──────────────────────────────────────────
        #
        # MNE's apply() always does:
        #   1. x_pre  = x_V / pre_whitener_          (pre-whiten: V → µV if pre_whitener_=1/scale)
        #   2. x_pca  = pca_components_ @ (x_pre - pca_mean_)
        #   3. sources = unmixing_matrix_ @ x_pca
        #   [reconstruction reverses in opposite order, multiplying by pre_whitener_ at the end]
        #
        # AMICA forward pass (µV-scale):
        #   sources = W @ sphere @ (x_µV - mean_µV)
        #           = W @ V @ diag(D⁻½) @ V.T @ (x_µV - mean_µV)
        #
        # Setting pre_whitener_ = 1/scale means x_pre = x_V * scale = x_µV.
        # Then with pca_components_ = V.T and pca_mean_ = mean_µV:
        #   x_pca  = V.T @ (x_µV - mean_µV)                             ✓
        #   unmixing_matrix_ = W @ V @ diag(D⁻½)  (no scale factor)     ✓
        #   mixing_matrix_   = diag(D^½) @ V.T @ A = (V*d_sqrt).T @ A   ✓
        #
        # Round-trip (no exclusions, full rank):
        #   pre_whitener_ * (V @ (V.T@(x_µV-mean_µV)) + mean_µV) = (1/scale)*x_µV = x_V  ✓

        n_ch             = V.shape[0]
        d_sqrt           = np.sqrt(d_vals)
        pre_whitener_    = np.full((n_ch, 1), 1.0 / scale)         # (n_ch, 1): V → µV
        pca_components_         = V.T.copy()                        # (n_comp, n_ch)
        pca_explained_variance_ = d_vals                            # (n_comp,) in µV²
        pca_mean_               = m.mean_.cpu().numpy()             # (n_ch,)   in µV

        # V * d_invsqrt: col j of V scaled by d_invsqrt[j]  =  V @ diag(D⁻½)
        unmixing_matrix_ = W @ (V * d_invsqrt)                     # (n_comp, n_comp)
        # inv(unmixing) = diag(D^½) @ V.T @ inv(W) = (V * d_sqrt).T @ A
        mixing_matrix_   = (V * d_sqrt).T @ A                      # (n_comp, n_comp)

        # ── Construct MNE ICA and populate fitted state ────────────────────
        ica = MNE_ICA(
            n_components = n_comp,
            method       = 'fastica',   # must pass validation; overwritten below
            max_iter     = m.max_iter,
        )
        ica.method = 'amica'            # honest label for repr / save

        ica.pre_whitener_           = pre_whitener_
        ica.pca_components_         = pca_components_
        ica.pca_explained_variance_ = pca_explained_variance_
        ica.pca_mean_               = pca_mean_
        ica.unmixing_matrix_        = unmixing_matrix_
        ica.mixing_matrix_          = mixing_matrix_
        ica.n_components_           = n_comp
        ica.n_pca_components_       = n_comp
        ica.n_samples_              = self._n_samples
        ica.ch_names                = list(self._ch_names)
        ica.info                    = self._inst_info
        ica.exclude                 = []
        ica.current_fit             = self._fit_type
        ica._ica_names              = [f'ICA{i:03d}' for i in range(n_comp)]
        ica.reject_                 = None
        ica.labels_                 = {}

        return ica
