"""Internal support module.

Conceptually, this module converts one session into area-level matrices,
coding-direction projections, and compact context metrics.

It exists as a separate unit so session-aware workflows can keep session I/O,
trial masking, and neural analysis as distinct layers.

It connects trial selection from `session_trials` to the higher-level
session-aware façade exposed through `session_aware`.
"""

from __future__ import annotations

from typing import Any, Dict, Mapping, Optional, Sequence, Tuple

import numpy as np
import pandas as pd

__all__ = [
    "build_cd_context",
    "compute_context_metrics",
    "extract_late_delay_matrix",
    "extract_trial_vectors_for_area",
    "extract_window_matrix",
    "get_area_units",
    "project_trials_on_cd",
    "safe_pearson",
    "safe_spearman",
    "sem",
]


def get_area_units(session: Mapping[str, Any], area_name: str, verbose: bool = False) -> pd.DataFrame:
    """Return the subset of `units_df` belonging to one target area."""

    units_df = session["units_df"]
    if "location" not in units_df.columns:
        if verbose:
            print(f"[units] {session.get('session_id', 'session')}: colonne 'location' absente dans units_df")
        return units_df.iloc[0:0].copy()

    area_norm = units_df["location"].astype(str).str.strip().str.lower()
    return units_df.loc[area_norm == str(area_name).strip().lower()].copy()


def extract_trial_vectors_for_area(
    session: Mapping[str, Any],
    area_name: str,
    trials: Sequence[int],
    t_start: float,
    t_end: float,
    align_to: str = "auditory_stim_time",
    verbose: bool = False,
) -> Tuple[np.ndarray, np.ndarray]:
    """Extract one area-level firing-rate matrix for a time window.

    Parameters
    ----------
    session:
        Session mapping containing at least `trials_df` and `units_df`.
    area_name:
        Area label used to filter the unit table.
    trials:
        Trial indices to evaluate.
    t_start, t_end:
        Window bounds in seconds relative to `align_to`.
    align_to:
        Trial-table column providing the alignment timestamp.

    Returns
    -------
    np.ndarray, np.ndarray
        Firing-rate matrix with shape `(valid_trials, neurons)` and the
        retained trial indices after alignment filtering.
    """

    trials_df = session["trials_df"]
    units_area = get_area_units(session, area_name, verbose=verbose)

    trial_indices = np.asarray(trials, dtype=int).reshape(-1)
    if trial_indices.size == 0 or units_area.empty:
        return np.empty((0, len(units_area)), dtype=float), np.array([], dtype=int)

    if align_to not in trials_df.columns:
        raise KeyError(f"Missing align_to column: {align_to}")

    align_times = pd.to_numeric(trials_df[align_to], errors="coerce").to_numpy(dtype=float)
    # Drop trials with invalid anchors so each output row still represents one
    # well-defined analysis window.
    valid_mask = np.isfinite(align_times[trial_indices])
    valid_trials = trial_indices[valid_mask]
    if verbose and np.any(~valid_mask):
        bad_n = int(np.sum(~valid_mask))
        print(
            f"[align] {session.get('session_id', 'session')} {area_name}: "
            f"{bad_n} trials ignorés car {align_to} est non fini"
        )
    if valid_trials.size == 0:
        return np.empty((0, len(units_area)), dtype=float), np.array([], dtype=int)

    duration = float(t_end - t_start)
    if duration <= 0:
        raise ValueError("t_end must be larger than t_start.")

    if "spike_times" not in units_area.columns:
        raise KeyError("Colonne 'spike_times' absente dans units_df")

    spike_arrays = []
    for spike_times in units_area["spike_times"].values:
        try:
            spike_arr = np.asarray(spike_times, dtype=float).reshape(-1)
        except Exception:
            spike_arr = np.array([], dtype=float)
        spike_arrays.append(spike_arr)

    X = np.zeros((valid_trials.size, len(spike_arrays)), dtype=float)
    for i_trial, trial_idx in enumerate(valid_trials):
        anchor = float(align_times[trial_idx])
        w0 = anchor + float(t_start)
        w1 = anchor + float(t_end)
        for i_unit, spike_arr in enumerate(spike_arrays):
            if spike_arr.size == 0:
                X[i_trial, i_unit] = 0.0
            else:
                # Convert spike counts to firing rate so windows with different
                # widths remain directly comparable.
                X[i_trial, i_unit] = np.sum((spike_arr >= w0) & (spike_arr < w1)) / duration

    return X, valid_trials


def extract_window_matrix(
    session: Mapping[str, Any],
    area_name: str,
    trial_indices: Sequence[int],
    align_to: str = "auditory_stim_time",
    window: Tuple[float, float] = (-0.2, 0.0),
    verbose: bool = False,
    return_valid_trials: bool = False,
):
    """Return an area-level matrix for one `(start, end)` analysis window.

    Parameters
    ----------
    session:
        Session mapping containing neural and trial metadata.
    area_name:
        Area label used to filter the unit table.
    trial_indices:
        Trial indices to evaluate.
    align_to:
        Trial-table column providing the alignment timestamp.
    window:
        `(start, end)` bounds in seconds relative to `align_to`.
    return_valid_trials:
        If True, also return the retained trial indices after alignment
        filtering.

    Returns
    -------
    np.ndarray or tuple[np.ndarray, np.ndarray]
        Area-level matrix, optionally paired with the retained trial indices.
    """

    X, valid_trials = extract_trial_vectors_for_area(
        session=session,
        area_name=area_name,
        trials=trial_indices,
        t_start=float(window[0]),
        t_end=float(window[1]),
        align_to=align_to,
        verbose=verbose,
    )
    if return_valid_trials:
        return X, valid_trials
    return X


def extract_late_delay_matrix(
    session: Mapping[str, Any],
    area_name: str,
    trial_indices: Sequence[int],
    align_to: str = "whisker_stim_time",
    window: Tuple[float, float] = (-0.2, 0.0),
    verbose: bool = False,
) -> np.ndarray:
    """Return the late-delay window matrix used by session-aware notebooks."""

    return extract_window_matrix(
        session=session,
        area_name=area_name,
        trial_indices=trial_indices,
        align_to=align_to,
        window=window,
        verbose=verbose,
        return_valid_trials=False,
    )


def build_cd_context(X_go: np.ndarray, X_nogo: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Build a simple context coding direction from Go and Nogo matrices.

    Parameters
    ----------
    X_go, X_nogo:
        Condition matrices with shape `(trials, neurons)`.

    Returns
    -------
    np.ndarray, np.ndarray, np.ndarray, np.ndarray
        `(cd, midpoint, mean_go, mean_nogo)`.
    """

    X_go = np.asarray(X_go, dtype=float)
    X_nogo = np.asarray(X_nogo, dtype=float)

    if X_go.ndim != 2 or X_nogo.ndim != 2:
        raise ValueError("X_go and X_nogo must be 2D arrays (trials, neurons).")
    if X_go.shape[1] != X_nogo.shape[1]:
        raise ValueError("X_go and X_nogo must have the same number of neurons.")
    if X_go.shape[0] == 0 or X_nogo.shape[0] == 0:
        raise ValueError("Need at least one Go and one Nogo trial to build CDcontext.")

    mean_go = np.nanmean(X_go, axis=0)
    mean_nogo = np.nanmean(X_nogo, axis=0)
    cd = mean_go - mean_nogo
    midpoint = 0.5 * (mean_go + mean_nogo)
    return cd, midpoint, mean_go, mean_nogo


def project_trials_on_cd(X_trials: np.ndarray, cd: np.ndarray, midpoint: Optional[np.ndarray] = None) -> np.ndarray:
    """Project trial vectors onto a unit-normalized coding direction.

    Parameters
    ----------
    X_trials:
        Trial matrix with shape `(trials, neurons)`.
    cd:
        Coding direction with shape `(neurons,)`.
    midpoint:
        Optional centering vector applied before projection.

    Returns
    -------
    np.ndarray
        One projected score per trial.
    """

    X_trials = np.asarray(X_trials, dtype=float)
    cd = np.asarray(cd, dtype=float).reshape(-1)

    if X_trials.ndim != 2:
        raise ValueError("X_trials must be 2D (trials, neurons).")
    if X_trials.shape[1] != cd.size:
        raise ValueError("X_trials and cd must share the same neuron dimension.")

    cd_norm = float(np.linalg.norm(cd))
    if (not np.isfinite(cd_norm)) or cd_norm == 0.0:
        return np.full(X_trials.shape[0], np.nan, dtype=float)

    cd_unit = cd / cd_norm
    centered = X_trials if midpoint is None else X_trials - np.asarray(midpoint, dtype=float).reshape(1, -1)
    return centered @ cd_unit


def compute_context_metrics(scores_go: np.ndarray, scores_nogo: np.ndarray) -> Dict[str, float]:
    """Compute simple separation metrics from Go and Nogo score vectors.

    Parameters
    ----------
    scores_go, scores_nogo:
        One-dimensional score vectors for the two conditions.

    Returns
    -------
    dict[str, float]
        Summary metrics including sample counts, central tendency, d-prime, and
        ROC AUC when both classes contain at least one finite score.
    """

    scores_go = np.asarray(scores_go, dtype=float).reshape(-1)
    scores_nogo = np.asarray(scores_nogo, dtype=float).reshape(-1)

    scores_go = scores_go[np.isfinite(scores_go)]
    scores_nogo = scores_nogo[np.isfinite(scores_nogo)]

    n_go = int(scores_go.size)
    n_nogo = int(scores_nogo.size)

    mean_go = float(np.nanmean(scores_go)) if n_go > 0 else np.nan
    mean_nogo = float(np.nanmean(scores_nogo)) if n_nogo > 0 else np.nan
    median_go = float(np.nanmedian(scores_go)) if n_go > 0 else np.nan
    median_nogo = float(np.nanmedian(scores_nogo)) if n_nogo > 0 else np.nan
    mean_diff = mean_go - mean_nogo if np.isfinite(mean_go) and np.isfinite(mean_nogo) else np.nan

    pooled_std = np.nan
    d_prime = np.nan
    if n_go >= 2 and n_nogo >= 2:
        # Use the pooled within-class variance so d-prime remains comparable
        # across sessions with different trial counts per condition.
        var_go = float(np.nanvar(scores_go, ddof=1))
        var_nogo = float(np.nanvar(scores_nogo, ddof=1))
        denom_df = n_go + n_nogo - 2
        if denom_df > 0:
            pooled_var = (((n_go - 1) * var_go) + ((n_nogo - 1) * var_nogo)) / denom_df
            if np.isfinite(pooled_var) and pooled_var > 0:
                pooled_std = float(np.sqrt(pooled_var))
                d_prime = float(mean_diff / pooled_std)

    auc = np.nan
    if n_go > 0 and n_nogo > 0:
        try:
            from sklearn.metrics import roc_auc_score

            y_true = np.concatenate([np.ones(n_go, dtype=int), np.zeros(n_nogo, dtype=int)])
            y_score = np.concatenate([scores_go, scores_nogo])
            auc = float(roc_auc_score(y_true, y_score))
        except Exception:
            # Keep metrics available even when ROC computation is not possible in
            # the current environment or the score vectors are degenerate.
            auc = np.nan

    return {
        "n_go": n_go,
        "n_nogo": n_nogo,
        "mean_go": mean_go,
        "mean_nogo": mean_nogo,
        "mean_diff": mean_diff,
        "pooled_std": pooled_std,
        "d_prime": d_prime,
        "auc": auc,
        "median_go": median_go,
        "median_nogo": median_nogo,
    }


def safe_pearson(x: np.ndarray, y: np.ndarray) -> float:
    """Return a NaN-safe Pearson correlation for two 1D vectors."""

    x = np.asarray(x, dtype=float).reshape(-1)
    y = np.asarray(y, dtype=float).reshape(-1)
    valid = np.isfinite(x) & np.isfinite(y)
    x = x[valid]
    y = y[valid]
    if x.size < 2:
        return np.nan
    if np.allclose(np.std(x), 0.0) or np.allclose(np.std(y), 0.0):
        return np.nan
    return float(np.corrcoef(x, y)[0, 1])


def safe_spearman(x: np.ndarray, y: np.ndarray) -> float:
    """Return a NaN-safe Spearman correlation for two 1D vectors."""

    x = pd.Series(np.asarray(x, dtype=float).reshape(-1))
    y = pd.Series(np.asarray(y, dtype=float).reshape(-1))
    valid = x.notna() & y.notna()
    if int(valid.sum()) < 2:
        return np.nan
    xr = x.loc[valid].rank(method="average").to_numpy(dtype=float)
    yr = y.loc[valid].rank(method="average").to_numpy(dtype=float)
    return safe_pearson(xr, yr)


def sem(values: np.ndarray) -> float:
    """Return a NaN-safe standard error of the mean for one vector."""

    values = np.asarray(values, dtype=float)
    values = values[np.isfinite(values)]
    if values.size <= 1:
        return np.nan
    return float(np.nanstd(values, ddof=1) / np.sqrt(values.size))