"""Public helper module.

Conceptually, this module assembles area-level PSTH matrices used by the
structured plotting notebooks.

It exists as a separate unit so condition filtering, cell masking, baseline
handling, and cell-wise aggregation stay consistent across multiple figures.

It connects canonical per-area entries to the compact `{concatsig,
window_centers}` structures expected by PSTH plotting code.
"""

from __future__ import annotations

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

import numpy as np

from .core import unwrap_scalar_obj
from .decoding import get_ccf_mask, get_celltype_mask, get_completion_mask, get_quiet_mask
from .figure_data import nanmean_no_warn, nanstd_no_warn


def mean_sem_over_cells(signal: Any, n_bins: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray]:
    """Return mean and SEM across cells for a `(bins, cells)` signal matrix."""
    arr = np.asarray(signal, dtype=np.float32)
    if arr.ndim != 2:
        if n_bins is None:
            raise ValueError("n_bins is required when signal is not 2D")
        shape = (int(n_bins),)
        return np.full(shape, np.nan, dtype=np.float32), np.full(shape, np.nan, dtype=np.float32)
    if arr.shape[1] == 0:
        shape = (arr.shape[0],)
        return np.full(shape, np.nan, dtype=np.float32), np.full(shape, np.nan, dtype=np.float32)

    mean_sig = nanmean_no_warn(arr, axis=1).astype(np.float32, copy=False)
    # Code_M plot_condition_psths.m / plot_layer_psths.m use std(rows,0,1,'omitnan')
    # (flag 0 = sample std, N-1) divided by sqrt of the total row count.
    sem_sig = (nanstd_no_warn(arr, axis=1, ddof=1) / np.sqrt(max(arr.shape[1], 1))).astype(np.float32, copy=False)
    return mean_sig, sem_sig


def classify_cortical_layers(layer_like: Any) -> np.ndarray:
    """Classify layer-like strings into manuscript layer groups.

    Parameters
    ----------
    layer_like:
        One-dimensional or nested payload containing Allen/CCF layer labels.

    Returns
    -------
    np.ndarray
        Object array of normalized layer labels using
        `{'supragranular', 'granular', 'infragranular'}`.
    """

    values = np.asarray(unwrap_scalar_obj(layer_like), dtype=object).reshape(-1)
    out: List[str] = []
    for raw in values:
        # Mirror Code_M classify_layers.m exactly: strip brackets/whitespace and
        # use bare-substring membership (contains(s,'5'), contains(s,'4'),
        # contains(s,'1')) plus strcmp(s,'6'). The previous narrowed forms
        # ("layer 5", " 5", ...) reclassified some Allen labels differently.
        token = str(unwrap_scalar_obj(raw)).replace("[", "").replace("]", "").strip()
        if ("6a" in token) or ("6b" in token) or (token == "6") or ("5" in token):
            out.append("infragranular")
        elif "4" in token:
            out.append("granular")
        elif ("2/3" in token) or ("1" in token):
            out.append("supragranular")
        else:
            # Keep unmatched labels in the deepest bucket so every neuron can
            # still contribute to the layerwise panels and tables.
            out.append("infragranular")
    return np.asarray(out, dtype=object)


def layer_display_name(layer_name: str) -> str:
    """Return the manuscript-facing display label for one normalized layer."""

    mapping = {
        "supragranular": "Supragranular",
        "granular": "Granular",
        "infragranular": "Infragranular",
    }
    key = str(layer_name).strip().lower()
    return mapping.get(key, key.title())


def collect_area_condition_signal(
    entries_by_area: Mapping[str, List[Mapping[str, Any]]],
    area_list: Optional[Mapping[str, Set[str]]],
    window_centers: np.ndarray,
    current_area: str,
    trial_type: int,
    lick_state: int,
    *,
    quiet_state: str,
    completion_state: str,
    cell_type: str,
    baseline_subtraction: bool,
    baseline_window: Tuple[float, float] = (0.95, 1.0),
    enable_ccf_filter: bool = True,
) -> np.ndarray:
    """Collect one concatenated `(bins, cells)` PSTH matrix for one condition.

    Parameters
    ----------
    entries_by_area:
        Canonical entries grouped by area.
    area_list:
        Optional area-to-CCF mapping used by the unit filter.
    window_centers:
        Shared PSTH time axis.
    current_area:
        Area to extract.
    trial_type, lick_state:
        Trial-state filters defining the requested condition.
    quiet_state, completion_state, cell_type:
        Trial and unit selection settings.
    baseline_subtraction:
        Whether to subtract a per-cell baseline before concatenation.
    baseline_window:
        Baseline window in seconds on `window_centers`.
    enable_ccf_filter:
        Whether CCF-based unit filtering is active.

    Returns
    -------
    np.ndarray
        Concatenated signal matrix with shape `(bins, cells)`.
    """
    concat_sp: List[np.ndarray] = []
    baseline_first_bin = int(np.argmin(np.abs(window_centers - float(baseline_window[0]))))
    baseline_last_bin = int(np.argmin(np.abs(window_centers - float(baseline_window[1]))))

    for entry in entries_by_area.get(current_area, []):
        trial = np.asarray(entry["trial"]).reshape(-1)
        lick = np.asarray(entry["lick"]).reshape(-1)
        ind_trial_type = trial == int(trial_type)
        ind_lick_state = lick == int(lick_state)
        curr_trial_ind = (
            get_quiet_mask(entry, quiet_state)
            & get_completion_mask(entry, completion_state)
            & ind_lick_state
            & ind_trial_type
        )

        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)) or (not np.any(curr_trial_ind)):
            continue

        curr_sp = np.asarray(entry["spike_counts"], dtype=np.float32)
        # Average across selected trials first, then keep only the requested
        # cell subset so concatenation stays cell-wise across sessions.
        curr_sp_trials = nanmean_no_warn(curr_sp[:, curr_trial_ind, :], axis=1).astype(np.float32, copy=False)
        curr_sp_trials_cells = curr_sp_trials[:, curr_mask]
        if curr_sp_trials_cells.size == 0:
            continue

        if baseline_subtraction:
            baseline_mean = np.repeat(
                nanmean_no_warn(
                    curr_sp_trials_cells[baseline_first_bin : baseline_last_bin + 1, :],
                    axis=0,
                ).reshape(1, -1),
                curr_sp_trials_cells.shape[0],
                axis=0,
            ).astype(np.float32, copy=False)
            curr_sp_trials_cells = curr_sp_trials_cells - baseline_mean

        concat_sp.append(curr_sp_trials_cells.astype(np.float32, copy=False))

    if not concat_sp:
        return np.zeros((window_centers.size, 0), dtype=np.float32)
    return np.concatenate(concat_sp, axis=1).astype(np.float32, copy=False)


def build_area_condition_psth_data(
    entries_by_area: Mapping[str, List[Mapping[str, Any]]],
    area_list: Optional[Mapping[str, Set[str]]],
    window_centers: np.ndarray,
    regionlist: Sequence[str],
    trial_types: Sequence[int],
    lick_states: Sequence[int],
    *,
    quiet_state: str,
    completion_state: str,
    cell_type: str,
    baseline_subtraction: bool,
    baseline_window: Tuple[float, float] = (-1.0, 0.0),
    enable_ccf_filter: bool = True,
) -> Dict[str, Dict[int, Dict[str, Any]]]:
    """Build `area -> condition -> {concatsig, window_centers}` mappings.

    Parameters
    ----------
    entries_by_area:
        Canonical entries grouped by area.
    area_list:
        Optional area-to-CCF mapping used by the unit filter.
    window_centers:
        Shared PSTH time axis.
    regionlist:
        Areas to materialize in the output mapping.
    trial_types, lick_states:
        Parallel condition definitions consumed pairwise.
    quiet_state, completion_state, cell_type:
        Trial and unit selection settings.
    baseline_subtraction:
        Whether to subtract a per-cell baseline before concatenation.
    baseline_window:
        Baseline window in seconds on `window_centers`.
    enable_ccf_filter:
        Whether CCF-based unit filtering is active.

    Returns
    -------
    dict[str, dict[int, dict[str, Any]]]
        Area-grouped condition payloads ready for PSTH plotting notebooks.
    """
    out: Dict[str, Dict[int, Dict[str, Any]]] = {}
    wc = np.asarray(window_centers, dtype=np.float32).reshape(-1)

    for area_name in regionlist:
        out[area_name] = {}
        for icond, (trial_type, lick_state) in enumerate(zip(trial_types, lick_states)):
            concatsig = collect_area_condition_signal(
                entries_by_area,
                area_list,
                wc,
                area_name,
                int(trial_type),
                int(lick_state),
                quiet_state=quiet_state,
                completion_state=completion_state,
                cell_type=cell_type,
                baseline_subtraction=baseline_subtraction,
                baseline_window=baseline_window,
                enable_ccf_filter=enable_ccf_filter,
            )
            out[area_name][icond] = {
                "concatsig": concatsig,
                "window_centers": wc,
            }
    return out


def build_celltype_area_condition_psth_data(
    entries_by_area: Mapping[str, List[Mapping[str, Any]]],
    area_list: Optional[Mapping[str, Set[str]]],
    window_centers: np.ndarray,
    regionlist: Sequence[str],
    cell_types: Sequence[str],
    trial_types: Sequence[int],
    lick_states: Sequence[int],
    *,
    quiet_state: str,
    completion_state: str,
    baseline_subtraction: bool,
    baseline_window: Tuple[float, float] = (-1.0, 0.0),
    enable_ccf_filter: bool = True,
) -> Dict[str, Dict[str, Dict[int, Dict[str, Any]]]]:
    """Build `cell_type -> area -> condition -> {concatsig, window_centers}`.

    Parameters
    ----------
    entries_by_area:
        Canonical entries grouped by area.
    area_list:
        Optional area-to-CCF mapping used by the unit filter.
    window_centers:
        Shared PSTH time axis.
    regionlist:
        Areas to materialize in the output mapping.
    cell_types:
        Cell-type filters to evaluate independently.
    trial_types, lick_states:
        Parallel condition definitions consumed pairwise.
    quiet_state, completion_state:
        Trial selection settings.
    baseline_subtraction:
        Whether to subtract a per-cell baseline before concatenation.
    baseline_window:
        Baseline window in seconds on `window_centers`.
    enable_ccf_filter:
        Whether CCF-based unit filtering is active.

    Returns
    -------
    dict[str, dict[str, dict[int, dict[str, Any]]]]
        Nested cell-type and area mapping ready for PSTH plotting notebooks.
    """
    out: Dict[str, Dict[str, Dict[int, Dict[str, Any]]]] = {}
    for current_cell_type in cell_types:
        out[current_cell_type] = build_area_condition_psth_data(
            entries_by_area,
            area_list,
            window_centers,
            regionlist,
            trial_types,
            lick_states,
            quiet_state=quiet_state,
            completion_state=completion_state,
            cell_type=current_cell_type,
            baseline_subtraction=baseline_subtraction,
            baseline_window=baseline_window,
            enable_ccf_filter=enable_ccf_filter,
        )
    return out