"""Internal support module.

Conceptually, this module groups the trial-mask and quiet-trial helpers used by
the session-aware analysis stack.

It exists as a separate unit so trial classification, quiet-trial recovery,
and boolean coercion stay reusable across session-level analyses.

It connects raw session trial tables and behavioral time series to the boolean
trial masks consumed by `session_context`.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Tuple

import numpy as np
import pandas as pd

from .behavior import QuietTrialParams, nwb_find_quiet_trial, psth_behavior, robust_mode
from .loading import load_nwb_behavior_timeseries_data, load_nwb_behavior_timeseries_timestamps
from .session_io import PathLike, _session_label

__all__ = [
    "get_completed_trials",
    "get_context_trial_masks",
    "get_go_hit_trials",
    "get_go_whisker_trials",
    "get_nogo_cr_trials",
    "get_nogo_whisker_trials",
    "get_quiet_trials",
    "get_session_quiet_trials",
    "to_bool_array",
]


def to_bool_array(values: Any) -> np.ndarray:
    """Convert a pandas-like column to a robust boolean NumPy array."""

    if isinstance(values, pd.Series):
        series = values.copy()
    else:
        series = pd.Series(values)

    if series.dtype == bool:
        return series.to_numpy(dtype=bool)

    if pd.api.types.is_numeric_dtype(series):
        return series.fillna(0).to_numpy(dtype=float) != 0

    norm = series.astype(str).str.strip().str.lower()
    return norm.isin({"1", "true", "t", "yes", "y"}).to_numpy(dtype=bool)


def get_completed_trials(trials_df: pd.DataFrame, verbose: bool = False) -> np.ndarray:
    """Return the completed-trial mask with a safe all-trials fallback."""

    if "early_lick" not in trials_df.columns:
        if verbose:
            print("[completed] colonne 'early_lick' absente -> fallback: tous les trials sont considérés completed")
        return np.ones(len(trials_df), dtype=bool)
    return ~to_bool_array(trials_df["early_lick"])


def get_quiet_trials(
    nwb: Any,
    trials_df: pd.DataFrame,
    session_path: PathLike,
    use_quiet_filter: bool = True,
    session_id: Optional[str] = None,
    verbose: bool = False,
) -> Tuple[np.ndarray, str]:
    """Return a quiet-trial mask and its provenance label.

    Parameters
    ----------
    nwb:
        Open NWB object for the current session.
    trials_df:
        Session trial table.
    session_path:
        Source NWB path, used for session-specific sampling-rate handling and
        readable fallback messages.
    use_quiet_filter:
        Whether quiet-trial filtering is enabled.
    session_id:
        Optional explicit session label used in verbose messages.
    verbose:
        Whether fallback decisions should be printed.

    Returns
    -------
    tuple[np.ndarray, str]
        Quiet-trial mask and a short provenance string describing how it was
        obtained.
    """

    n_trials = len(trials_df)
    if not use_quiet_filter:
        return np.ones(n_trials, dtype=bool), "quiet filter off"

    all_cols = list(trials_df.columns)
    q_wh_col = "quiet_trial_whisker_speed" if "quiet_trial_whisker_speed" in all_cols else None
    q_jaw_col = "quiet_trial_jaw_movement" if "quiet_trial_jaw_movement" in all_cols else None

    # Reuse precomputed quiet-trial columns when they already exist in the
    # trials table; this avoids rebuilding the behavioral PSTHs.
    if q_wh_col is not None and q_jaw_col is not None:
        quiet_mask = to_bool_array(trials_df[q_wh_col]) & to_bool_array(trials_df[q_jaw_col])
        return quiet_mask, f"{q_jaw_col} & {q_wh_col}"

    if "start_time" not in trials_df.columns:
        if verbose:
            print(f"[quiet] {_session_label(session_id, session_path)}: colonne 'start_time' absente -> fallback all trials")
        return np.ones(n_trials, dtype=bool), "fallback all trials (missing start_time)"

    try:
        anchor_times = pd.to_numeric(trials_df["start_time"], errors="coerce").to_numpy(dtype=np.float32)
        frame_rate = 100.0 if "PG019" in Path(session_path).name else 200.0

        jaw_coord = np.asarray(
            load_nwb_behavior_timeseries_data(nwb, "BehavioralTimeSeries", "Jaw_Coordinate"),
            dtype=np.float32,
        )
        whisker_angle = np.asarray(
            load_nwb_behavior_timeseries_data(nwb, "BehavioralTimeSeries", "C2Whisker_Angle"),
            dtype=np.float32,
        ).reshape(-1)
        video_timestamp = np.asarray(
            load_nwb_behavior_timeseries_timestamps(nwb, "BehavioralTimeSeries", "C2Whisker_Angle"),
            dtype=np.float32,
        )

        if whisker_angle.size <= 1:
            whisker_speed = np.zeros_like(whisker_angle, dtype=np.float32)
        else:
            diff_abs = np.abs(np.diff(whisker_angle.astype(np.float32), n=1))
            whisker_speed = np.concatenate([diff_abs[:1], diff_abs]).astype(np.float32)

        jaw_x = jaw_coord[:, 0]
        jaw_y = jaw_coord[:, 1]
        jaw_mode_x = float(robust_mode(jaw_x))
        jaw_mode_y = float(robust_mode(jaw_y))
        jaw_movement = np.sqrt((jaw_y - jaw_mode_y) ** 2 + (jaw_x - jaw_mode_x) ** 2).astype(np.float32)

        movement_signal = {
            "whisker_speed": whisker_speed,
            "jaw_movement": jaw_movement,
        }

        psth_beh, window_centers = psth_behavior(
            movement_signal,
            video_timestamp,
            anchor_times,
            pre_time=-1.0,
            post_time=2.0,
            window_size=0.01,
            window_step=0.01,
            fs=frame_rate,
        )

        behavior_table = dict(psth_beh)
        behavior_table["trial_timestamps"] = np.asarray(window_centers["whisker_speed"], dtype=np.float32)
        quiet_params = QuietTrialParams(
            prewhisk_window=(0.8, 1.0),
            baseline_window=(-1.0, 0.0),
            movement_signals=("whisker_speed", "jaw_movement"),
            selection_method="mad_all",
        )
        behavior_table = nwb_find_quiet_trial(behavior_table, quiet_params)

        q_wh = np.asarray(behavior_table["quiet_trial_whisker_speed"]).reshape(-1).astype(bool)
        q_jaw = np.asarray(behavior_table["quiet_trial_jaw_movement"]).reshape(-1).astype(bool)

        # Pad or trim the recovered quiet masks so they stay aligned to the
        # session trial table even if one source is slightly shorter.
        if q_wh.size != n_trials:
            q_wh_fixed = np.zeros(n_trials, dtype=bool)
            q_wh_fixed[: min(n_trials, q_wh.size)] = q_wh[: min(n_trials, q_wh.size)]
            q_wh = q_wh_fixed

        if q_jaw.size != n_trials:
            q_jaw_fixed = np.zeros(n_trials, dtype=bool)
            q_jaw_fixed[: min(n_trials, q_jaw.size)] = q_jaw[: min(n_trials, q_jaw.size)]
            q_jaw = q_jaw_fixed

        return q_wh & q_jaw, "computed quiet_trial_jaw_movement & quiet_trial_whisker_speed"

    except Exception as exc:
        if verbose:
            print(f"[quiet] {_session_label(session_id, session_path)}: fallback all trials ({type(exc).__name__}: {exc})")
        return np.ones(n_trials, dtype=bool), f"fallback all trials ({type(exc).__name__})"


def get_session_quiet_trials(
    session: Mapping[str, Any],
    use_quiet_filter: bool = True,
    verbose: bool = False,
) -> Tuple[np.ndarray, str]:
    """Wrapper around `get_quiet_trials` for standard session mappings."""

    return get_quiet_trials(
        nwb=session["nwb"],
        trials_df=session["trials_df"],
        session_path=session["session_path"],
        use_quiet_filter=use_quiet_filter,
        session_id=str(session.get("session_id", "")),
        verbose=verbose,
    )


def get_go_whisker_trials(trials_df: pd.DataFrame) -> np.ndarray:
    """Return the mask for Go-tone whisker trials."""

    if "trial_type" in trials_df.columns:
        trial_type_num = pd.to_numeric(trials_df["trial_type"], errors="coerce").to_numpy()
        return trial_type_num == 1

    if {"context", "whisker_stim"}.issubset(trials_df.columns):
        context = trials_df["context"].astype(str).str.strip().str.lower().to_numpy()
        whisker = to_bool_array(trials_df["whisker_stim"])
        return (context == "go_tone") & whisker

    raise KeyError("Impossible de définir les trials Go whisker.")


def get_nogo_whisker_trials(trials_df: pd.DataFrame) -> np.ndarray:
    """Return the mask for Nogo-tone whisker trials."""

    if "trial_type" in trials_df.columns:
        trial_type_num = pd.to_numeric(trials_df["trial_type"], errors="coerce").to_numpy()
        return trial_type_num == 3

    if {"context", "whisker_stim"}.issubset(trials_df.columns):
        context = trials_df["context"].astype(str).str.strip().str.lower().to_numpy()
        whisker = to_bool_array(trials_df["whisker_stim"])
        return (context == "nogo_tone") & whisker

    raise KeyError("Impossible de définir les trials Nogo whisker.")


def get_go_hit_trials(trials_df: pd.DataFrame) -> np.ndarray:
    """Return the mask for Go-tone whisker Hit trials."""

    go_mask = get_go_whisker_trials(trials_df)
    if "lick_flag" not in trials_df.columns:
        return go_mask
    return go_mask & to_bool_array(trials_df["lick_flag"])


def get_nogo_cr_trials(trials_df: pd.DataFrame) -> np.ndarray:
    """Return the mask for Nogo-tone whisker Correct Rejection trials."""

    nogo_mask = get_nogo_whisker_trials(trials_df)
    if "lick_flag" not in trials_df.columns:
        return nogo_mask
    return nogo_mask & (~to_bool_array(trials_df["lick_flag"]))


def get_context_trial_masks(trials_df: pd.DataFrame) -> Dict[str, np.ndarray]:
    """Return the canonical Go Hit and Nogo CR masks for context analyses."""

    missing = [col for col in ["trial_type", "lick_flag"] if col not in trials_df.columns]
    if missing:
        raise KeyError(f"Colonnes requises absentes pour définir Go Hit / Nogo CR: {missing}")

    trial_type_num = pd.to_numeric(trials_df["trial_type"], errors="coerce").to_numpy()
    lick_flag = to_bool_array(trials_df["lick_flag"])
    return {
        "go_whisker_hit": (trial_type_num == 1) & lick_flag,
        "nogo_whisker_cr": (trial_type_num == 3) & (~lick_flag),
    }