"""Public helper module.

Conceptually, this module keeps the lightweight helpers used by structured
coding-direction figure notebooks.

It exists as a separate unit so notebooks can stay linear while the payload
loading, session-aligned condition summaries, and recurrent plotting logic live
in reusable functions.

It connects saved coding-direction projections to the temporal and response-
window summaries rendered in Figure 6B-style notebooks.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Mapping, Sequence
import warnings

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import kruskal

from .core import unwrap_scalar_obj
from .figure_data import align_to_len, mean_sem_over_axis, nanmean_no_warn, nanstd_no_warn, to_dict_like
from .figure_decoding import window_index_bounds
from .figure_stats import add_prettify_like_pvalues, pairwise_kruskal_lsd_pvals, signrank_p_value
from .figure_style import add_line_legend, add_time_event_lines, style_plot_axis

__all__ = [
    "load_coding_direction_npz",
    "plot_cd_context_figure",
    "summarize_cd_area_projection",
    "summarize_spontlick_area_projection",
]


def _nanmean_no_warning(values: Any, axis: int | None = None) -> np.ndarray:
    """Return `np.nanmean` while silencing all-NaN-slice runtime warnings."""

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)
        return np.nanmean(np.asarray(values, dtype=float), axis=axis)


def load_coding_direction_npz(path: str | Path) -> tuple[dict[str, Any], np.ndarray]:
    """Load one saved coding-direction payload and its time-bin vector."""

    path = Path(path)
    with np.load(path, allow_pickle=True) as data:
        if "coding_direction_matrix" not in data:
            raise KeyError(f"{path.name} missing key 'coding_direction_matrix'")
        if "windowCenters" not in data:
            raise KeyError(f"{path.name} missing key 'windowCenters'")

        coding_direction_matrix = to_dict_like(np.asarray(data["coding_direction_matrix"]).reshape(-1)[0])
        window_centers = np.asarray(data["windowCenters"], dtype=float).reshape(-1)

    return coding_direction_matrix, window_centers


def summarize_cd_area_projection(
    coding_direction_matrix: Mapping[str, Any],
    *,
    area_name: str,
    window_centers: Any,
    condition_configs: Sequence[Mapping[str, Any]],
    response_window: tuple[float, float],
    baseline_subtraction: bool = False,
    baseline_bin_count: int = 100,
) -> dict[str, Any]:
    """Summarize one area's context-CD trajectories across behavioral conditions.

    The returned condition matrices keep one column per stored session and fill
    missing condition/session combinations with `NaN`, preserving paired session
    alignment for downstream signed-rank tests.
    """

    centers = np.asarray(window_centers, dtype=float).reshape(-1)
    if centers.size == 0:
        raise ValueError("window_centers cannot be empty")

    sessions_raw = coding_direction_matrix.get(area_name, np.array([], dtype=object))
    sessions_arr = np.asarray(unwrap_scalar_obj(sessions_raw), dtype=object).reshape(-1)
    n_sessions = int(sessions_arr.size)
    response_first, response_last = window_index_bounds(centers, response_window)

    condition_summaries: list[dict[str, Any]] = []
    p_values = np.full(len(condition_configs), np.nan, dtype=float)
    ref_values = np.full(n_sessions, np.nan, dtype=float)
    bottom_tops: list[float] = []

    for cond_index, condition in enumerate(condition_configs):
        context_concat = np.full((centers.size, n_sessions), np.nan, dtype=float)

        for session_index, session_raw in enumerate(sessions_arr):
            try:
                session = to_dict_like(session_raw)
                current_context = np.asarray(unwrap_scalar_obj(session["Contextproj"]), dtype=float)
                index_dict = to_dict_like(session["index"])
            except Exception:
                continue

            if current_context.ndim != 2 or current_context.shape[0] != centers.size:
                continue

            n_trials = int(current_context.shape[1])
            lick = align_to_len(index_dict.get("lick", []), n_trials, bool)
            trial = align_to_len(index_dict.get("trial", []), n_trials, float)
            quiet = align_to_len(index_dict.get("Quiet", []), n_trials, bool)
            completed = align_to_len(index_dict.get("completed_trial", []), n_trials, bool)

            trial_mask = (
                (trial == float(condition["trial"]))
                & (lick.astype(int) == int(condition["lick"]))
                & quiet
                & completed
            )
            if not np.any(trial_mask):
                continue

            curr = current_context[:, trial_mask]
            if curr.size == 0:
                continue

            if baseline_subtraction:
                base_stop = min(int(baseline_bin_count), curr.shape[0])
                if base_stop > 0:
                    baseline = np.mean(curr[:base_stop, :], axis=0, keepdims=True)
                    curr = curr - baseline

            context_concat[:, session_index] = _nanmean_no_warning(curr, axis=1)

        mean_trace, sem_trace = mean_sem_over_axis(context_concat, axis=1)
        response_values = _nanmean_no_warning(context_concat[response_first : response_last + 1, :], axis=0)
        response_mean = float(nanmean_no_warn(response_values)) if response_values.size else float("nan")
        finite_n = int(np.sum(np.isfinite(response_values)))
        response_sem = (
            float(nanstd_no_warn(response_values, ddof=1) / np.sqrt(max(finite_n, 1)))
            if response_values.size
            else float("nan")
        )

        top_candidates: list[float] = []
        if np.isfinite(response_mean):
            top_candidates.append(response_mean)
        if np.isfinite(response_sem):
            top_candidates.append(response_mean + response_sem)
        finite_response = response_values[np.isfinite(response_values)]
        if finite_response.size > 0:
            top_candidates.append(float(np.nanmax(finite_response)))
        if top_candidates:
            bottom_tops.append(float(np.nanmax(top_candidates)))

        if cond_index == 0:
            ref_values = response_values.copy()
            p_values[cond_index] = 1.0
        else:
            p_values[cond_index] = signrank_p_value(ref_values, response_values)

        condition_summaries.append(
            {
                "label": str(condition.get("label", f"Condition {cond_index + 1}")),
                "context_concat": context_concat,
                "mean_trace": mean_trace,
                "sem_trace": sem_trace,
                "response_values": response_values,
                "response_mean": response_mean,
                "response_sem": response_sem,
            }
        )

    finite_bottom_tops = [value for value in bottom_tops if np.isfinite(value)]
    bracket_y_position = min(2.15, max(finite_bottom_tops) + 0.18) if finite_bottom_tops else 1.55

    return {
        "area_name": str(area_name),
        "n_sessions": n_sessions,
        "condition_summaries": condition_summaries,
        "p_values": p_values,
        "bracket_y_position": float(bracket_y_position),
    }


def _aligned_cd_session_payload(
    session_raw: Any,
    *,
    projection_key: str,
) -> dict[str, np.ndarray]:
    """Return one coding-direction session aligned on a shared trial axis.

    Parameters
    ----------
    session_raw:
        One stored session entry from a coding-direction payload.
    projection_key:
        Projection matrix to extract from the session, such as `Contextproj`
        or `lickproj`.

    Returns
    -------
    dict[str, np.ndarray]
        Session-aligned projection and index vectors with matching trial count.
    """

    session = to_dict_like(session_raw)
    projection = np.asarray(unwrap_scalar_obj(session[projection_key]), dtype=float)
    if projection.ndim != 2:
        raise ValueError(f"{projection_key} must be a 2D `(time, trials)` array")

    index_dict = to_dict_like(session["index"])
    n_trials = int(projection.shape[1])
    return {
        "projection": projection[:, :n_trials],
        "lick": align_to_len(index_dict.get("lick", []), n_trials, bool),
        "trial": align_to_len(index_dict.get("trial", []), n_trials, float),
        "quiet": align_to_len(index_dict.get("Quiet", []), n_trials, bool),
        "completed": align_to_len(index_dict.get("completed_trial", []), n_trials, bool),
    }


def _condition_trial_mask(payload: Mapping[str, np.ndarray], condition: Mapping[str, Any]) -> np.ndarray:
    """Return the boolean trial mask for one condition definition."""

    trial = np.asarray(payload["trial"], dtype=float).reshape(-1)
    lick = np.asarray(payload["lick"], dtype=bool).reshape(-1)
    quiet = np.asarray(payload["quiet"], dtype=bool).reshape(-1)
    completed = np.asarray(payload["completed"], dtype=bool).reshape(-1)

    if "selector" in condition:
        cond_mask = np.asarray(condition["selector"](trial, lick), dtype=bool).reshape(-1)
    else:
        cond_mask = (trial == float(condition["trial"])) & (lick.astype(int) == int(condition["lick"]))

    return cond_mask & quiet & completed


def summarize_spontlick_area_projection(
    coding_direction_matrix: Mapping[str, Any],
    *,
    area_name: str,
    window_centers: Any,
    condition_configs: Sequence[Mapping[str, Any]],
    response_window: tuple[float, float],
    min_num_trials: int = 2,
    baseline_subtraction: bool = False,
    baseline_bin_count: int = 200,
) -> dict[str, Any]:
    """Summarize one area's spontaneous-lick coding-direction responses.

    Parameters
    ----------
    coding_direction_matrix:
        Saved coding-direction payload keyed by cortical area.
    area_name:
        Area to summarize.
    window_centers:
        Time axis associated with the projection matrices.
    condition_configs:
        Sequence of condition definitions. Each condition may provide either a
        `selector(trial, lick)` callable or explicit `trial` and `lick` keys.
    response_window:
        Inclusive time window used for the bottom-row response summary.
    min_num_trials:
        Minimum number of valid quiet and completed trials required to keep a
        session for every requested condition.
    baseline_subtraction:
        Whether to subtract one per-trial baseline before averaging traces.
    baseline_bin_count:
        Number of leading time bins used for baseline subtraction.

    Returns
    -------
    dict[str, Any]
        Area-level traces, paired response summaries, and condition-level
        statistics ready for Figure 8B-style plotting.
    """

    centers = np.asarray(window_centers, dtype=float).reshape(-1)
    if centers.size == 0:
        raise ValueError("window_centers cannot be empty")

    sessions_raw = coding_direction_matrix.get(area_name, np.array([], dtype=object))
    sessions_arr = np.asarray(unwrap_scalar_obj(sessions_raw), dtype=object).reshape(-1)
    n_sessions = int(sessions_arr.size)
    response_first, response_last = window_index_bounds(centers, response_window)

    aligned_sessions: list[dict[str, np.ndarray] | None] = []
    drop_mask = np.zeros((n_sessions, len(condition_configs)), dtype=bool)

    for session_index, session_raw in enumerate(sessions_arr):
        try:
            payload = _aligned_cd_session_payload(session_raw, projection_key="lickproj")
            if payload["projection"].shape[0] != centers.size:
                raise ValueError("Projection length does not match `window_centers`")
        except Exception:
            aligned_sessions.append(None)
            drop_mask[session_index, :] = True
            continue

        aligned_sessions.append(payload)
        for cond_index, condition in enumerate(condition_configs):
            trial_mask = _condition_trial_mask(payload, condition)
            if int(np.sum(trial_mask)) < int(min_num_trials):
                drop_mask[session_index, cond_index] = True

    session_keep = np.where(~np.any(drop_mask, axis=1))[0]
    condition_summaries: list[dict[str, Any]] = []
    p_vs_reference = np.full(len(condition_configs), np.nan, dtype=float)
    mean_values = np.full((session_keep.size, len(condition_configs)), np.nan, dtype=float)
    pair_inputs: list[np.ndarray] = []

    for cond_index, condition in enumerate(condition_configs):
        lick_concat = np.full((centers.size, n_sessions), np.nan, dtype=float)

        # Keep one column per stored session so the paired comparisons stay aligned.
        for session_index, payload in enumerate(aligned_sessions):
            if payload is None:
                continue

            trial_mask = _condition_trial_mask(payload, condition)
            if not np.any(trial_mask):
                continue

            current_projection = payload["projection"][:, trial_mask]
            if current_projection.size == 0:
                continue

            if baseline_subtraction:
                base_stop = min(int(baseline_bin_count), current_projection.shape[0])
                if base_stop > 0:
                    baseline = np.mean(current_projection[:base_stop, :], axis=0, keepdims=True)
                    current_projection = current_projection - baseline

            lick_concat[:, session_index] = _nanmean_no_warning(current_projection, axis=1)

        signal2plot = (
            lick_concat[:, session_keep]
            if session_keep.size > 0
            else np.full((centers.size, 0), np.nan, dtype=float)
        )
        mean_trace, sem_trace = mean_sem_over_axis(signal2plot, axis=1)

        if session_keep.size == 0:
            response_values = np.array([], dtype=float)
        else:
            response_values = _nanmean_no_warning(lick_concat[response_first : response_last + 1, session_keep], axis=0)

        response_mean = float(nanmean_no_warn(response_values)) if response_values.size else float("nan")
        finite_n = int(np.sum(np.isfinite(response_values)))
        response_sem = (
            float(nanstd_no_warn(response_values, ddof=1) / np.sqrt(max(finite_n, 1)))
            if response_values.size
            else float("nan")
        )

        if response_values.size > 0:
            pair_inputs.append(np.column_stack([response_values, np.full(response_values.shape, cond_index + 1, dtype=float)]))

        if session_keep.size > 0 and response_values.size == session_keep.size:
            mean_values[:, cond_index] = response_values
            p_vs_reference[cond_index] = 1.0 if cond_index == 0 else signrank_p_value(mean_values[:, 0], mean_values[:, cond_index])

        condition_summaries.append(
            {
                "label": str(condition.get("label", f"Condition {cond_index + 1}")),
                "mean_trace": mean_trace,
                "sem_trace": sem_trace,
                "response_values": response_values,
                "response_mean": response_mean,
                "response_sem": response_sem,
            }
        )

    if pair_inputs:
        pair_matrix = np.vstack(pair_inputs)
        groups = [pair_matrix[pair_matrix[:, 1] == group_id, 0] for group_id in np.unique(pair_matrix[:, 1])]
        groups = [group[np.isfinite(group)] for group in groups if np.sum(np.isfinite(group)) > 0]
        if len(groups) >= 2:
            try:
                kw_stat, kw_p = kruskal(*groups)
            except Exception:
                kw_stat, kw_p = np.nan, np.nan
        else:
            kw_stat, kw_p = np.nan, np.nan

        pair_left, pair_right, pair_pvals = pairwise_kruskal_lsd_pvals(pair_matrix[:, 0], pair_matrix[:, 1])
    else:
        kw_stat, kw_p = np.nan, np.nan
        pair_left = np.array([], dtype=float)
        pair_right = np.array([], dtype=float)
        pair_pvals = np.array([], dtype=float)

    return {
        "area_name": str(area_name),
        "n_sessions": n_sessions,
        "n_session_keep": int(session_keep.size),
        "drop_mask": drop_mask,
        "condition_summaries": condition_summaries,
        "stats": {
            "kw_stat": kw_stat,
            "kw_p": kw_p,
            "pair_left": pair_left,
            "pair_right": pair_right,
            "pair_pvals": pair_pvals,
            "p_vs_reference": p_vs_reference,
        },
    }


def plot_cd_context_figure(
    area_summaries: Mapping[str, Mapping[str, Any]],
    *,
    area_names: Sequence[str],
    window_centers: Any,
    condition_colors: Any,
    time_ticks: Sequence[float],
    time_tick_labels: Sequence[str],
    top_y_limits: tuple[float, float],
    bottom_y_limits: tuple[float, float],
    figure_size: tuple[float, float] = (16.0, 7.0),
    dpi: int = 150,
    legend_labels: Sequence[str] | None = None,
    bottom_xtick_labels: Sequence[str] | None = None,
    significance_kwargs: Mapping[str, Any] | None = None,
) -> tuple[Any, np.ndarray]:
    """Render the two-row context-CD summary figure across all requested areas."""

    centers = np.asarray(window_centers, dtype=float).reshape(-1)
    colors = np.asarray(condition_colors, dtype=float)
    n_conditions = int(colors.shape[0])
    bottom_labels = list(bottom_xtick_labels or legend_labels or [f"C{i + 1}" for i in range(n_conditions)])
    sig_kwargs = dict(significance_kwargs or {})

    plt.rcParams.update(
        {
            "font.size": 10,
            "axes.titlesize": 12,
            "axes.labelsize": 11,
            "xtick.labelsize": 9,
            "ytick.labelsize": 9,
        }
    )

    fig, axs = plt.subplots(2, len(area_names), figsize=figure_size, dpi=dpi)
    axs = np.asarray(axs).reshape(2, len(area_names))

    for area_index, area_name in enumerate(area_names):
        ax_top = axs[0, area_index]
        ax_bottom = axs[1, area_index]
        summary = area_summaries.get(area_name, None)
        if not summary:
            ax_top.set_xlim([centers[0], centers[-1]])
            ax_top.set_ylim(list(top_y_limits))
            ax_bottom.set_xlim([0, n_conditions + 1])
            ax_bottom.set_ylim(list(bottom_y_limits))
            ax_top.set_title(str(area_name), fontweight="bold", pad=8)
            continue

        for cond_index, condition_summary in enumerate(summary["condition_summaries"]):
            mean_trace = np.asarray(condition_summary["mean_trace"], dtype=float)
            sem_trace = np.asarray(condition_summary["sem_trace"], dtype=float)
            response_values = np.asarray(condition_summary["response_values"], dtype=float)
            response_mean = float(condition_summary["response_mean"])
            response_sem = float(condition_summary["response_sem"])
            x_pos = cond_index + 1

            ax_top.fill_between(
                centers,
                mean_trace - sem_trace,
                mean_trace + sem_trace,
                color=colors[cond_index, :],
                alpha=0.30,
                linewidth=0,
            )
            ax_top.plot(centers, mean_trace, color=colors[cond_index, :], linewidth=1.0)

            ax_bottom.bar(
                [x_pos],
                [response_mean],
                color=colors[cond_index, :],
                width=0.68,
                edgecolor="black",
                linewidth=0.6,
            )
            ax_bottom.errorbar(
                [x_pos],
                [response_mean],
                [response_sem],
                fmt="-k",
                capsize=3,
                linewidth=1.0,
                markersize=4,
            )

            finite_values = response_values[np.isfinite(response_values)]
            if finite_values.size > 0:
                ax_bottom.plot(
                    np.full(finite_values.shape, x_pos + 0.2),
                    finite_values,
                    "o",
                    markersize=4,
                    markerfacecolor="none",
                    color="k",
                    alpha=0.55,
                )

        ax_top.set_xlim([centers[0], centers[-1]])
        ax_top.set_ylim(list(top_y_limits))
        ax_bottom.set_xlim([0, n_conditions + 1])
        ax_bottom.set_ylim(list(bottom_y_limits))

        if n_conditions > 1:
            add_prettify_like_pvalues(
                ax_bottom,
                [1] * (n_conditions - 1),
                list(range(2, n_conditions + 1)),
                np.asarray(summary["p_values"], dtype=float)[1:],
                y_position=float(summary["bracket_y_position"]),
                **sig_kwargs,
            )

        add_time_event_lines(ax_top, xs=(0.0, 1.0), color="k", linewidth=0.9)
        ax_top.axhline(0.0, color="k", linewidth=0.9)
        style_plot_axis(ax_top, facecolor="#f8f8f8", grid_axis="y", grid_color="0.82", grid_linewidth=0.6, grid_alpha=0.65, spine_width=1.0, tick_width=1.0, tick_length=3.5)
        style_plot_axis(ax_bottom, facecolor="#f8f8f8", grid_axis="y", grid_color="0.82", grid_linewidth=0.6, grid_alpha=0.65, spine_width=1.0, tick_width=1.0, tick_length=3.5)
        ax_bottom.set_xticks(list(range(1, n_conditions + 1)))
        ax_bottom.set_xticklabels(bottom_labels, rotation=35, ha="right")
        ax_top.set_title(str(area_name), fontweight="bold", pad=8)

    for ax in axs[0, :]:
        ax.set_xticks(list(time_ticks))
        ax.set_xticklabels(list(time_tick_labels))
        ax.set_xlabel("Time (s)")

    axs[0, 0].set_ylabel("CD Context (a.u)")
    axs[1, 0].set_ylabel("Mean delay CD Context")

    if legend_labels is not None:
        add_line_legend(
            fig,
            legend_labels,
            [colors[i, :] for i in range(min(len(legend_labels), n_conditions))],
            loc="upper center",
            bbox_to_anchor=(0.5, 1.02),
            ncol=max(len(legend_labels), 1),
            frameon=False,
            columnspacing=1.2,
            handlelength=2.2,
            linewidth=2.2,
        )

    fig.tight_layout(rect=[0.02, 0.03, 0.98, 0.95])
    return fig, axs