"""Public helper module.

Conceptually, this module assembles ROC and modulation payloads into the
summary arrays used by structured figure notebooks.

It exists as a separate unit so ROC-entry loading, area-wise aggregation, and
class summaries remain reusable across related figures.

It connects saved ROC/modulation payloads to the compact vectors and PSTH
matrices consumed by the plotting code.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List, Mapping, Optional, Sequence

import numpy as np

from .core import to_dict_like
from .decoding import get_ccf_mask, get_celltype_mask, get_completion_mask, get_quiet_mask
from .figure_data import mean_sem_over_axis


def unwrap_roc_entry(x: Any) -> Optional[Dict[str, Any]]:
    """Return a plain dict ROC entry or `None` when the payload is unsupported."""
    try:
        return to_dict_like(x)
    except Exception:
        return None


def load_roc_npz(path: Path, key: str = "roc_mat") -> np.ndarray:
    """Load a flat object-array ROC payload from an NPZ file."""
    path = Path(path)
    with np.load(path, allow_pickle=True) as d:
        if key not in d:
            raise KeyError(f"{path.name} missing key {key!r}")
        roc = np.asarray(d[key]).reshape(-1)
    return roc


def load_modulation_npz(path: Path) -> Dict[str, Dict[str, np.ndarray]]:
    """Load modulation payloads into a plain `period -> {'roc_mat': array}` mapping."""
    path = Path(path)
    with np.load(path, allow_pickle=True) as d:
        if "modulation" not in d:
            raise KeyError(f"{path.name} missing key 'modulation'")
        raw = d["modulation"]

    top = to_dict_like(raw)
    out: Dict[str, Dict[str, np.ndarray]] = {}
    for period, payload in top.items():
        payload_dict = to_dict_like(payload)
        if "roc_mat" not in payload_dict:
            raise KeyError(f"modulation[{period!r}] missing key 'roc_mat'")
        out[str(period)] = {"roc_mat": np.asarray(payload_dict["roc_mat"]).reshape(-1)}
    return out


def collect_area_roc_classification(
    entries: Sequence[Mapping[str, Any]],
    roc_arr: np.ndarray,
    current_area: str,
    area_list: Optional[Mapping[str, set[str]]],
    *,
    cell_type: str = "All",
    enable_ccf_filter: bool = True,
) -> Dict[str, Any]:
    """Collect concatenated ROC-classification vectors for one area.

    Parameters
    ----------
    entries:
        Canonical entries iterated in the same order as `roc_arr`.
    roc_arr:
        Flat ROC entry array loaded from disk.
    current_area:
        Area to extract.
    area_list:
        Optional area-to-CCF mapping used by the unit filter.
    cell_type:
        Cell-type filter applied inside each entry.
    enable_ccf_filter:
        Whether CCF-based unit filtering is active.

    Returns
    -------
    dict[str, Any]
        Concatenated discrimination-index, p-value, and delta-rate vectors plus
        the number of probes contributing to the area.
    """
    concat_pval: List[np.ndarray] = []
    concat_disc: List[np.ndarray] = []
    concat_diff: List[np.ndarray] = []
    n_probes = 0

    for fallback_idx, entry in enumerate(entries):
        if str(entry["probe_location"]) != str(current_area):
            continue
        n_probes += 1

        entry_idx = int(entry.get("entry_idx", fallback_idx))
        if entry_idx >= len(roc_arr):
            continue

        curr_roc = unwrap_roc_entry(roc_arr[entry_idx])
        if curr_roc is None:
            continue

        curr_mask = get_celltype_mask(entry, cell_type) & get_ccf_mask(
            entry,
            current_area,
            area_list=area_list,
            enable_ccf_filter=enable_ccf_filter,
        )
        if not np.any(curr_mask):
            continue

        roc_diff = np.asarray(curr_roc["diff_fr"], dtype=np.float32).reshape(-1)
        disc = np.asarray(curr_roc["discrimination_index"], dtype=np.float32).reshape(-1)
        pval = np.asarray(curr_roc["pvalue"], dtype=np.float32).reshape(-1)

        n = min(curr_mask.size, roc_diff.size, disc.size, pval.size)
        if n == 0:
            continue

        cm = np.asarray(curr_mask[:n], dtype=bool)
        if not np.any(cm):
            continue

        concat_diff.append(roc_diff[:n][cm])
        concat_disc.append(disc[:n][cm])
        concat_pval.append(pval[:n][cm])

    if concat_disc:
        disc_all = np.concatenate(concat_disc).astype(np.float32, copy=False)
        pval_all = np.concatenate(concat_pval).astype(np.float32, copy=False)
        diff_all = np.concatenate(concat_diff).astype(np.float32, copy=False)
    else:
        disc_all = np.asarray([], dtype=np.float32)
        pval_all = np.asarray([], dtype=np.float32)
        diff_all = np.asarray([], dtype=np.float32)

    return {
        "disc": disc_all,
        "pval": pval_all,
        "diff": diff_all,
        "n_probes": int(n_probes),
    }


def summarize_roc_classes(
    disc: Any,
    pval: Any,
    diff: Any,
) -> Dict[str, Any]:
    """Summarize positive, negative, and non-modulated ROC classes."""
    disc_arr = np.asarray(disc, dtype=np.float32).reshape(-1)
    pval_arr = np.asarray(pval, dtype=np.float32).reshape(-1)
    diff_arr = np.asarray(diff, dtype=np.float32).reshape(-1)

    n = min(disc_arr.size, pval_arr.size, diff_arr.size)
    disc_arr = disc_arr[:n]
    pval_arr = pval_arr[:n]
    diff_arr = diff_arr[:n]

    ind_pos = (disc_arr > 0) & (pval_arr < 0.05)
    ind_neg = (disc_arr < 0) & (pval_arr < 0.05)
    ind_non = ~(ind_pos | ind_neg)

    counts = np.array([np.sum(ind_pos), np.sum(ind_neg), np.sum(ind_non)], dtype=np.float32)
    if disc_arr.size > 0:
        fractions = counts / float(disc_arr.size)
    else:
        fractions = np.full(3, np.nan, dtype=np.float32)

    means, sems = [], []
    for mask in (ind_pos, ind_neg, ind_non):
        vals = diff_arr[mask]
        if vals.size == 0:
            means.append(np.nan)
            sems.append(np.nan)
            continue
        mean_v, sem_v = mean_sem_over_axis(vals, axis=0)
        means.append(float(np.asarray(mean_v)))
        sems.append(float(np.asarray(sem_v)))

    return {
        "counts": counts,
        "fractions": np.asarray(fractions, dtype=np.float32),
        "delta_means": np.asarray(means, dtype=np.float32),
        "delta_errors": np.asarray(sems, dtype=np.float32),
        "ind_pos": ind_pos,
        "ind_neg": ind_neg,
        "ind_non": ind_non,
    }


def collect_area_roc_psth_data(
    entries: Sequence[Mapping[str, Any]],
    roc_arr: np.ndarray,
    current_area: str,
    *,
    cell_type: str = "All",
    ind_window: int = 0,
    quiet_state: str = "Quiet_(jaw & whisker)",
    completion_state: str = "completed_trials",
    pos_trial_type: int = 1,
    pos_lick_state: int = 1,
    neg_trial_type: int = 3,
    neg_lick_state: int = 0,
    fallback_window_centers: Optional[np.ndarray] = None,
) -> Dict[str, Any]:
    """Collect PSTH matrices and ROC vectors for one area.

    Parameters
    ----------
    entries:
        Canonical entries iterated in the same order as `roc_arr`.
    roc_arr:
        Flat ROC entry array loaded from disk.
    current_area:
        Area to extract.
    cell_type:
        Cell-type filter applied inside each entry.
    ind_window:
        ROC window index used when the ROC payload is two-dimensional.
    quiet_state, completion_state:
        Trial-selection settings.
    pos_trial_type, pos_lick_state:
        Positive-condition trial and lick filters.
    neg_trial_type, neg_lick_state:
        Negative-condition trial and lick filters.
    fallback_window_centers:
        Optional time axis used when no area entries contribute data.

    Returns
    -------
    dict[str, Any]
        Area-level ROC vectors and PSTH matrices ready for plotting.
    """
    concat_pval: List[np.ndarray] = []
    concat_roc: List[np.ndarray] = []
    concat_pos: List[np.ndarray] = []
    concat_neg: List[np.ndarray] = []
    time2plot = None

    for fallback_idx, entry in enumerate(entries):
        if str(entry["probe_location"]) != str(current_area):
            continue

        entry_idx = int(entry.get("entry_idx", fallback_idx))
        if entry_idx >= len(roc_arr):
            continue

        roc_entry = unwrap_roc_entry(roc_arr[entry_idx])
        if roc_entry is None:
            continue

        curr_mask = get_celltype_mask(entry, cell_type)

        current_roc = np.asarray(roc_entry["discrimination_index"], dtype=np.float32)
        current_pvalue = np.asarray(roc_entry["pvalue"], dtype=np.float32)
        if current_roc.ndim == 2:
            current_roc = current_roc[:, int(ind_window)]
        else:
            current_roc = current_roc.reshape(-1)
        if current_pvalue.ndim == 2:
            current_pvalue = current_pvalue[:, int(ind_window)]
        else:
            current_pvalue = current_pvalue.reshape(-1)

        spike_counts = np.asarray(entry["spike_counts"], dtype=np.float32)
        n_units = min(curr_mask.size, current_roc.size, current_pvalue.size, int(spike_counts.shape[2]))
        if n_units == 0:
            continue
        curr_mask = np.asarray(curr_mask[:n_units], dtype=bool)
        if not np.any(curr_mask):
            continue

        trial = np.asarray(entry["trial"]).reshape(-1)
        lick = np.asarray(entry["lick"], dtype=bool).reshape(-1)
        qind = get_quiet_mask(entry, quiet_state)
        comp = get_completion_mask(entry, completion_state)
        n_trials = min(trial.size, lick.size, qind.size, comp.size, spike_counts.shape[1])
        if n_trials == 0:
            continue

        trial = trial[:n_trials]
        lick = lick[:n_trials]
        qind = qind[:n_trials]
        comp = comp[:n_trials]
        spike_counts = spike_counts[:, :n_trials, :n_units]

        pos_ind = qind & comp & (trial == int(pos_trial_type)) & (lick.astype(int) == int(pos_lick_state))
        neg_ind = qind & comp & (trial == int(neg_trial_type)) & (lick.astype(int) == int(neg_lick_state))

        if not np.any(pos_ind) and not np.any(neg_ind):
            continue

        current_sc = spike_counts[:, :, curr_mask]
        # Collapse each condition to one `(bins, selected_cells)` mean matrix so
        # cells can be concatenated across probes within the same area.
        if np.any(pos_ind):
            pos_mean = np.nanmean(current_sc[:, pos_ind, :], axis=1).astype(np.float32, copy=False)
            if pos_mean.ndim == 1:
                pos_mean = pos_mean.reshape(-1, 1)
        else:
            pos_mean = np.full((current_sc.shape[0], int(np.sum(curr_mask))), np.nan, dtype=np.float32)
        if np.any(neg_ind):
            neg_mean = np.nanmean(current_sc[:, neg_ind, :], axis=1).astype(np.float32, copy=False)
            if neg_mean.ndim == 1:
                neg_mean = neg_mean.reshape(-1, 1)
        else:
            neg_mean = np.full((current_sc.shape[0], int(np.sum(curr_mask))), np.nan, dtype=np.float32)

        concat_pval.append(current_pvalue[:n_units][curr_mask])
        concat_roc.append(current_roc[:n_units][curr_mask])
        concat_pos.append(pos_mean)
        concat_neg.append(neg_mean)
        time2plot = np.asarray(entry["trial_timestamps"], dtype=np.float32).reshape(-1)

    if concat_roc:
        return {
            "time2plot": time2plot if time2plot is not None else np.asarray(fallback_window_centers, dtype=np.float32),
            "concat_roc": np.concatenate(concat_roc).astype(np.float32, copy=False),
            "concat_pval": np.concatenate(concat_pval).astype(np.float32, copy=False),
            "concat_pos": np.concatenate(concat_pos, axis=1).astype(np.float32, copy=False),
            "concat_neg": np.concatenate(concat_neg, axis=1).astype(np.float32, copy=False),
        }

    empty_ts = np.asarray(fallback_window_centers, dtype=np.float32).reshape(-1)
    return {
        "time2plot": empty_ts,
        "concat_roc": np.asarray([], dtype=np.float32),
        "concat_pval": np.asarray([], dtype=np.float32),
        "concat_pos": np.zeros((empty_ts.size, 0), dtype=np.float32),
        "concat_neg": np.zeros((empty_ts.size, 0), dtype=np.float32),
    }