"""Public helper module.

Conceptually, this module collects lightweight statistical transforms and
sampling utilities used across decoding, subspace, and figure-analysis code.

It exists as a separate unit so rank-based metrics, balancing utilities, and
array normalization remain consistent across higher-level pipelines.

It connects trial matrices and summary vectors to the scalar statistics needed
by decoding and plotting modules.
"""

from __future__ import annotations

from typing import Any, Optional, Tuple

import numpy as np

__all__ = [
    "angle_btw_nd_vectors",
    "auc_from_scores_labels",
    "corrcoef_with_pvalues",
    "downsample_balance",
    "holdout_split",
    "independent_group_stats",
    "selectivity_index_calculation_np",
    "zscore_cols",
    "zscore_matlab_cols",
]


def auc_from_scores_labels(
    scores: np.ndarray,
    labels: np.ndarray,
    pos_label: Any = 1,
) -> float:
    """Compute AUC from score ranks using a tie-safe Mann-Whitney formulation.

    Parameters
    ----------
    scores:
        Score vector whose larger values indicate stronger evidence for the
        positive class.
    labels:
        One-dimensional label vector aligned to `scores`.
    pos_label:
        Label treated as the positive class.

    Returns
    -------
    float
        Area under the ROC curve, or `nan` when one class is missing.
    """

    scores = np.asarray(scores, dtype=np.float64).reshape(-1)
    labels = np.asarray(labels).reshape(-1)
    if scores.size != labels.size:
        raise ValueError("scores and labels must have the same length")

    pos = labels == pos_label
    n_pos = int(np.sum(pos))
    n_neg = int(labels.size - n_pos)
    if n_pos == 0 or n_neg == 0:
        return float("nan")

    order = np.argsort(scores, kind="mergesort")
    s_sorted = scores[order]
    ranks = np.empty(scores.size, dtype=np.float64)

    # Assign average ranks within tie blocks so the resulting AUC matches the
    # standard Mann-Whitney interpretation.
    i = 0
    while i < scores.size:
        j = i + 1
        while j < scores.size and s_sorted[j] == s_sorted[i]:
            j += 1
        avg_rank = (i + 1 + j) / 2.0
        ranks[order[i:j]] = avg_rank
        i = j

    sum_pos = float(np.sum(ranks[pos]))
    auc = (sum_pos - (n_pos * (n_pos + 1) / 2.0)) / float(n_pos * n_neg)
    return float(auc)



def selectivity_index_calculation_np(
    data: np.ndarray,
    labels: np.ndarray,
    method: str = "permut",
    permutations: int = 200,
    nboot: int = 200,
    rng_obj: Optional[np.random.Generator] = None,
    pos_label: Any = 1,
) -> Tuple[np.float32, np.float32, np.float32]:
    """Compute a ROC-based selectivity index with permutation p-value.

    Parameters
    ----------
    data:
        One-dimensional response vector.
    labels:
        Class labels aligned to `data`.
    method:
        Selectivity-estimation strategy. Only `"permut"` is supported.
    permutations:
        Number of label permutations used for the p-value estimate.
    nboot:
        Unused legacy placeholder retained for call compatibility.
    rng_obj:
        Optional random generator used for permutation sampling.
    pos_label:
        Label treated as the positive class.

    Returns
    -------
    tuple[np.float32, np.float32, np.float32]
        Discrimination index, permutation p-value, and raw AUC.
    """

    del nboot

    if method != "permut":
        raise ValueError(f"Unsupported ROC method in Python code: {method}")

    # MATLAB selectivity_index_calculation.m works in double precision and the
    # discrimination index is exported as Selectivity_index; keep float64 so
    # float32 rounding does not create spurious ties that shift the rank-AUC.
    data = np.asarray(data, dtype=np.float64).reshape(-1)
    labels = np.asarray(labels).reshape(-1)
    if data.size != labels.size:
        raise ValueError("data and labels must have the same length")

    auc = auc_from_scores_labels(data, labels, pos_label=pos_label)
    if np.isnan(auc):
        return np.float64(np.nan), np.float64(np.nan), np.float64(np.nan)

    discrimination_index = np.float64((auc - 0.5) * 2.0)

    idx_rand = np.empty(int(permutations), dtype=np.float64)
    for i in range(int(permutations)):
        perm = np.random.permutation(labels.size) if rng_obj is None else rng_obj.permutation(labels.size)
        auc_rand = auc_from_scores_labels(data, labels[perm], pos_label=pos_label)
        idx_rand[i] = np.float64((auc_rand - 0.5) * 2.0)

    p_value = np.float64(np.mean(np.abs(discrimination_index) <= np.abs(idx_rand)))
    return discrimination_index, p_value, np.float64(auc)



def angle_btw_nd_vectors(vec1: np.ndarray, vec2: Optional[np.ndarray] = None) -> np.ndarray:
    """Return the pairwise angle matrix between one or two vector sets."""

    v1 = np.asarray(vec1, dtype=np.float64)
    if v1.ndim == 1:
        v1 = v1.reshape(1, -1)

    if vec2 is None:
        v2 = v1
    else:
        v2 = np.asarray(vec2, dtype=np.float64)
        if v2.ndim == 1:
            v2 = v2.reshape(1, -1)

    dot = v1 @ v2.T
    v1_norm = np.sqrt(np.sum(v1 * v1, axis=1))
    v2_norm = np.sqrt(np.sum(v2 * v2, axis=1))
    denom = np.outer(v1_norm, v2_norm)

    with np.errstate(invalid="ignore", divide="ignore"):
        cos_theta = dot / denom
    cos_theta = np.clip(cos_theta, -1.0, 1.0)
    return np.abs(np.arccos(cos_theta))



def corrcoef_with_pvalues(X: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """Return a correlation matrix with two-sided p-values per entry.

    Parameters
    ----------
    X:
        Observation matrix with shape `(observations, variables)`.

    Returns
    -------
    tuple[np.ndarray, np.ndarray]
        Correlation matrix and matching p-value matrix.
    """

    from scipy import stats

    x = np.asarray(X, dtype=np.float64)
    with np.errstate(invalid="ignore", divide="ignore"):
        R = np.corrcoef(x, rowvar=False)
    R = np.asarray(R, dtype=np.float64)
    if R.ndim == 0:
        R = R.reshape(1, 1)
    elif R.ndim == 1:
        R = np.atleast_2d(R)

    n_obs = x.shape[0] if x.ndim >= 2 else x.size
    dof = int(n_obs - 2)
    if dof <= 0:
        P = np.full_like(R, np.nan, dtype=np.float64)
        return R, P

    with np.errstate(divide="ignore", invalid="ignore"):
        t_stat = R * np.sqrt(dof / (1.0 - R * R))
        P = 2.0 * stats.t.sf(np.abs(t_stat), dof)

    P[np.isnan(R)] = np.nan
    return R, P



def holdout_split(
    ids: np.ndarray,
    holdout: float = 0.7,
    rng_obj: Optional[np.random.Generator] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """Split ids into train and test subsets using one random holdout.

    Parameters
    ----------
    ids:
        One-dimensional id vector.
    holdout:
        Fraction of ids assigned to the test subset.
    rng_obj:
        Optional random generator controlling the permutation.

    Returns
    -------
    tuple[np.ndarray, np.ndarray]
        Train ids followed by test ids.
    """

    ids = np.asarray(ids, dtype=int).reshape(-1)
    n = ids.size
    if n == 0:
        return np.array([], dtype=int), np.array([], dtype=int)

    n_test = int(np.floor(n * holdout))
    if n_test <= 0:
        return ids.copy(), np.array([], dtype=int)
    if n_test >= n:
        n_test = n - 1

    perm = rng_obj.permutation(n) if rng_obj is not None else np.random.permutation(n)
    test_pos = perm[:n_test]
    train_mask = np.ones(n, dtype=bool)
    train_mask[test_pos] = False

    train_ids = ids[train_mask]
    test_ids = ids[~train_mask]
    return train_ids, test_ids


def independent_group_stats(data1: np.ndarray, data2: np.ndarray) -> dict[str, float | int]:
    """Return a lightweight two-group summary with t-test and Cohen's d.

    Parameters
    ----------
    data1, data2:
        One-dimensional numeric samples to compare.

    Returns
    -------
    dict[str, float | int]
        Summary dictionary containing means, sample sizes, pooled-effect size,
        and a two-sided independent-samples t-test when SciPy is available.
    """

    x1 = np.asarray(data1, dtype=np.float64).reshape(-1)
    x2 = np.asarray(data2, dtype=np.float64).reshape(-1)
    x1 = x1[np.isfinite(x1)]
    x2 = x2[np.isfinite(x2)]

    n1 = int(x1.size)
    n2 = int(x2.size)
    if n1 == 0 or n2 == 0:
        return {
            "t_stat": float("nan"),
            "p_value": float("nan"),
            "cohens_d": float("nan"),
            "mean1": float("nan"),
            "mean2": float("nan"),
            "n1": n1,
            "n2": n2,
        }

    mean1 = float(np.mean(x1))
    mean2 = float(np.mean(x2))
    var1 = float(np.var(x1, ddof=1)) if n1 > 1 else 0.0
    var2 = float(np.var(x2, ddof=1)) if n2 > 1 else 0.0
    pooled_denom = int(n1 + n2 - 2)
    if pooled_denom > 0:
        pooled_std = np.sqrt((((n1 - 1) * var1) + ((n2 - 1) * var2)) / pooled_denom)
    else:
        pooled_std = float("nan")
    if np.isfinite(pooled_std) and pooled_std > 0:
        cohens_d = float((mean1 - mean2) / pooled_std)
    else:
        cohens_d = 0.0

    try:
        from scipy import stats

        t_stat, p_value = stats.ttest_ind(x1, x2, equal_var=True, nan_policy="omit")
        t_stat = float(t_stat)
        p_value = float(p_value)
    except Exception:
        t_stat = float("nan")
        p_value = float("nan")

    return {
        "t_stat": t_stat,
        "p_value": p_value,
        "cohens_d": cohens_d,
        "mean1": mean1,
        "mean2": mean2,
        "n1": n1,
        "n2": n2,
    }



def zscore_cols(X: np.ndarray) -> np.ndarray:
    """Apply column-wise z-scoring with sample standard deviation and NaN safety.

    Parameters
    ----------
    X:
        Two-dimensional feature matrix.

    Returns
    -------
    np.ndarray
        Float32 z-scored matrix with non-finite values replaced by zero.
    """

    X = np.asarray(X, dtype=np.float64)
    finite = np.isfinite(X)
    Xf = np.where(finite, X, 0.0)

    counts = finite.sum(axis=0).astype(np.int64)
    mu = np.zeros((X.shape[1],), dtype=np.float64)
    np.divide(Xf.sum(axis=0), counts, out=mu, where=counts > 0)

    xc = np.where(finite, X - mu[np.newaxis, :], 0.0)
    # Use sample variance column by column, while still returning zeros for
    # columns that never contain a finite observation.
    denom = np.maximum(counts - 1, 1)

    var = np.ones((X.shape[1],), dtype=np.float64)
    np.divide((xc * xc).sum(axis=0), denom, out=var, where=counts > 0)

    sd = np.sqrt(var)
    sd[~np.isfinite(sd)] = 1.0
    sd[sd == 0] = 1.0

    Z = xc / sd[np.newaxis, :]
    Z[:, counts == 0] = 0.0
    # MATLAB zscore returns double; this feeds covariance/PCA/subspace math, so
    # keep float64 to match Code_M precision.
    return np.nan_to_num(Z, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float64, copy=False)



def zscore_matlab_cols(X: np.ndarray) -> np.ndarray:
    """Backward-compatible alias for the column-wise z-score helper."""

    return zscore_cols(X)



def downsample_balance(
    X1: np.ndarray,
    X2: np.ndarray,
    rng_obj: Optional[np.random.Generator] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """Balance two trial matrices by random downsampling.

    Parameters
    ----------
    X1, X2:
        Trial matrices with shape `(trials, features)`.
    rng_obj:
        Optional random generator controlling the downsampling permutation.

    Returns
    -------
    tuple[np.ndarray, np.ndarray]
        Balanced feature matrix and the corresponding `+1/-1` label vector.
    """

    X1 = np.asarray(X1, dtype=np.float32)
    X2 = np.asarray(X2, dtype=np.float32)

    n1, n2 = X1.shape[0], X2.shape[0]
    if n1 == n2:
        X = np.vstack([X1, X2])
        y = np.concatenate([np.ones(n1), -np.ones(n2)])
        return X, y

    if rng_obj is None:
        rng_obj = np.random.default_rng()

    if n1 > n2:
        idx = rng_obj.permutation(n1)[:n2]
        X1 = X1[idx, :]
    else:
        idx = rng_obj.permutation(n2)[:n1]
        X2 = X2[idx, :]

    X = np.vstack([X1, X2])
    y = np.concatenate([np.ones(X1.shape[0]), -np.ones(X2.shape[0])])
    return X, y