from __future__ import annotations

from typing import Any

import mne
import numpy as np

from ..base import InverseOperator, SolverMeta
from .base_beamformer import BaseBeamformer
from .utils import (
    _infer_mvpure_n_sources_and_rank,
    _mvpure_projected_lcmv_weights_from_inv_cov,
    _select_top_diverse_indices,
    build_covariance_candidates,
)


class SolverESMVMVPURE(BaseBeamformer):
    """MV-PURE extension of ESMV with a projected multisource LCMV core.

    The vanilla ESMV core uses single-source LCMV weights followed by an
    eigenspace projection. This variant replaces the single-source LCMV stage
    with an MV-PURE projected multisource LCMV filter on a selected source set.
    """

    meta = SolverMeta(
        slug="esmv-mvpure",
        full_name="Eigenspace Minimum Variance (MV-PURE)",
        category="Beamformers",
        description=(
            "ESMV variant where the underlying MV filter is replaced by a "
            "reduced-rank MV-PURE multisource LCMV stage."
        ),
        references=[
            "Jonmohamadi, Y., Poudel, G., Innes, C., Weiss, D., Krueger, R., & Jones, R. "
            "(2014). Comparison of beamformers for EEG source signal reconstruction. "
            "Biomedical Signal Processing and Control, 14, 175-188.",
            "Jurkowska, J., Dreszer, J., Lewandowska, M., & Piotrowski, T. (2025). "
            "Multi-Source Neural Activity Indices and Spatial Filters for EEG/MEG "
            "Inverse Problem: An Extension to MNE-Python. arXiv preprint arXiv:2509.14118.",
        ],
    )
    R_VALUE_EXPONENTS = (-16.0, 1.0)

    def __init__(
        self,
        name: str = "ESMV MV-PURE Beamformer",
        *,
        mvp_n_sources: int | str = "auto",
        mvp_rank: int | str = "auto",
        mvp_max_sources: int = 4,
        spectrum_threshold: float = 0.05,
        selection_corr_threshold: float = 0.97,
        source_indices: list[int] | np.ndarray | None = None,
        reduce_rank: bool = True,
        rank: str = "auto",
        **kwargs,
    ):
        self.name = name
        self.mvp_n_sources = mvp_n_sources
        self.mvp_rank = mvp_rank
        self.mvp_max_sources = int(mvp_max_sources)
        self.spectrum_threshold = float(spectrum_threshold)
        self.selection_corr_threshold = float(selection_corr_threshold)
        self.source_indices = (
            None if source_indices is None else np.asarray(source_indices, dtype=int)
        )
        super().__init__(reduce_rank=reduce_rank, rank=rank, **kwargs)

    def _infer_l0_and_rank(self, C: np.ndarray) -> tuple[int, int]:
        return _infer_mvpure_n_sources_and_rank(
            C,
            mvp_n_sources=self.mvp_n_sources,
            mvp_rank=self.mvp_rank,
            mvp_max_sources=self.mvp_max_sources,
            spectrum_threshold=self.spectrum_threshold,
        )

    def make_inverse_operator(  # type: ignore[override]
        self,
        forward: Any,
        mne_obj: Any,
        *args: Any,
        alpha: str | float = "auto",
        weight_norm: bool = False,
        noise_cov: mne.Covariance | None = None,
        cov_reg: str = "oas",
        cov_reg_beta: float = 0.05,
        cov_reg_cond_target: float = 1e4,
        **kwargs: Any,
    ) -> Any:
        self.weight_norm = bool(weight_norm)
        super().make_inverse_operator(forward, mne_obj, *args, alpha=alpha, **kwargs)
        wf = self.prepare_whitened_forward(noise_cov)
        data = self.unpack_data_obj(mne_obj)

        leadfield = wf.G_white
        leadfield_norm = np.linalg.norm(leadfield, axis=0, keepdims=True)
        leadfield_norm = np.where(leadfield_norm > 0, leadfield_norm, 1.0)
        leadfield /= leadfield_norm
        n_chans, n_dipoles = leadfield.shape
        leadfield_similarity = np.abs(leadfield.T @ leadfield)

        y = wf.sensor_transform @ data
        y -= y.mean(axis=1, keepdims=True)
        I = np.identity(n_chans)
        C = self.data_covariance(y, center=False, ddof=1)
        C_sub = self.select_signal_subspace(C)
        cov_mats, self.alphas, cov_meta = build_covariance_candidates(
            C=C,
            I=I,
            alpha=self.alpha,
            get_alphas_fn=self.get_alphas,
            n_samples=int(y.shape[1]),
            cov_reg=cov_reg,
            cov_reg_beta=float(cov_reg_beta),
            cov_reg_cond_target=float(cov_reg_cond_target),
        )
        if "oas_shrinkage" in cov_meta:
            self._cov_reg_oas_shrinkage = float(cov_meta["oas_shrinkage"])

        inverse_operators = []
        selected_sources_per_alpha: list[np.ndarray] = []
        selected_rank_per_alpha: list[int] = []
        for cov_mat in cov_mats:
            C_inv = self.robust_inverse(cov_mat)

            if self.source_indices is not None:
                sel = np.unique(self.source_indices)
                if sel.size == 0:
                    raise ValueError("source_indices must be non-empty when provided")
                l0 = int(sel.size)
                r = (
                    int(self.mvp_rank)
                    if isinstance(self.mvp_rank, (int, np.integer))
                    else l0
                )
                r = max(1, min(r, l0))
            else:
                l0, r = self._infer_l0_and_rank(C)

                # Rank sources by vanilla single-source LCMV output power.
                upper = C_inv @ leadfield
                lower = np.einsum("ij,jk,ki->i", leadfield.T, C_inv, leadfield)
                lower = np.maximum(lower, 1e-20)
                W_vanilla = upper / lower
                q_hat = W_vanilla.T @ y
                scores = np.mean(q_hat**2, axis=1)
                scores = np.nan_to_num(scores, nan=0.0, posinf=0.0, neginf=0.0)

                sel = _select_top_diverse_indices(
                    scores,
                    leadfield_similarity,
                    l0,
                    corr_threshold=self.selection_corr_threshold,
                )
                sel = np.unique(sel)
                l0 = int(sel.size)
                r = max(1, min(r, l0))

            selected_sources_per_alpha.append(sel.copy())
            selected_rank_per_alpha.append(int(r))

            H0 = leadfield[:, sel]
            W_sel = _mvpure_projected_lcmv_weights_from_inv_cov(
                C_inv, H0, rank=r, cond_threshold=1e12, regularization=1e-12
            )  # (l0, n_chans)

            K_full = np.zeros((n_dipoles, n_chans), dtype=W_sel.dtype)
            K_full[sel, :] = W_sel
            K_esmv = K_full @ C_sub

            if self.weight_norm:
                row_norm = np.linalg.norm(K_esmv, axis=1)
                nz = row_norm > 0
                K_esmv[nz, :] = (K_esmv[nz, :].T / row_norm[nz]).T

            inverse_operators.append(K_esmv @ wf.sensor_transform)

        self.inverse_operators = [
            InverseOperator(inverse_operator, self.name)
            for inverse_operator in inverse_operators
        ]
        # Expose selection details for debugging/benchmarks.
        self.selected_sources_ = selected_sources_per_alpha
        self.selected_rank_ = selected_rank_per_alpha
        return self
