"""Public helper module.

Conceptually, this module assembles the session-wise summary values and paired
plots used by the optogenetic figure notebooks.

It exists as a separate unit so session metric extraction, multiple-comparison
correction, and paired-summary plotting stay synchronized across related
figures.

It connects session-level optogenetic payloads to the per-area summary panels
rendered in the structured notebooks.
"""

from __future__ import annotations

from collections.abc import Callable, Mapping, Sequence
from typing import Any

import matplotlib.pyplot as plt
import numpy as np

from .core import to_session_id_str
from .figure_data import fiber_matches, nanmean_no_warn, nanstd_no_warn
from .figure_stats import add_significance_simple, add_significance_stacked, bh_fdr, signrank_p_value
from .figure_style import style_plot_axis


def compute_area_window_session_metrics(
    optomat: Sequence[Mapping[str, Any]],
    *,
    region_list: Sequence[str],
    window_list: Sequence[str],
    trialtype_list: Sequence[int],
    metric_fn: Callable[[Mapping[str, Any], str, int], float],
) -> tuple[np.ndarray, dict[str, np.ndarray], dict[str, list[str]]]:
    """Compute per-session metrics by area, window, and trial type.

    Parameters
    ----------
    optomat:
        Session-level optogenetic payloads.
    region_list:
        Areas to evaluate.
    window_list:
        Ordered analysis windows. The first window is treated as the reference
        condition for within-session comparisons.
    trialtype_list:
        Trial-type codes evaluated independently.
    metric_fn:
        Callback returning one scalar session metric for `(session, window,
        trial_type)`.

    Returns
    -------
    tuple[np.ndarray, dict[str, np.ndarray], dict[str, list[str]]]
        Raw p-value tensor, area-indexed session metric arrays, and the
        session identifiers contributing to each area.
    """
    n_areas = len(region_list)
    n_trials = len(trialtype_list)
    n_windows = len(window_list)

    pvalue_matrix = np.full((n_areas, n_trials, n_windows - 1), np.nan, dtype=float)
    performance_by_area: dict[str, np.ndarray] = {}
    session_ids_by_area: dict[str, list[str]] = {}

    for ind_area, current_area in enumerate(region_list):
        fiberlist = [i for i, sess in enumerate(optomat) if fiber_matches(sess.get("fiber_location", ""), current_area)]
        area_sessions = [to_session_id_str(optomat[i].get("session_id", f"idx_{i}")) for i in fiberlist]

        for ind_trialtype, current_trial_type in enumerate(trialtype_list):
            performance = np.full((n_windows, len(fiberlist)), np.nan, dtype=float)
            pvals = np.full(n_windows, np.nan, dtype=float)

            # Evaluate the full window series for one area/trial-type block so
            # later plotting code can reuse both the metrics and the p-values.
            for ind_window, current_window in enumerate(window_list):
                for sess_idx, ind_session in enumerate(fiberlist):
                    performance[ind_window, sess_idx] = float(metric_fn(optomat[ind_session], current_window, int(current_trial_type)))

                if ind_window > 0:
                    pvals[ind_window] = signrank_p_value(performance[0, :], performance[ind_window, :])

            pvalue_matrix[ind_area, ind_trialtype, :] = pvals[1:]
            if ind_trialtype == 0:
                performance_by_area[current_area] = performance
                session_ids_by_area[current_area] = area_sessions

    return pvalue_matrix, performance_by_area, session_ids_by_area


def apply_fdr_correction(pvalue_matrix: np.ndarray, correction: str) -> np.ndarray:
    """Apply the requested BH-FDR strategy to a p-value tensor."""
    pvalue_matrix = np.asarray(pvalue_matrix, dtype=float)
    n_areas, n_trials, n_windows_minus_1 = pvalue_matrix.shape

    if correction == "area":
        out = np.full_like(pvalue_matrix, np.nan, dtype=float)
        for ind_area in range(n_areas):
            p_area = pvalue_matrix[ind_area, :, :].reshape(-1)
            out[ind_area, :, :] = bh_fdr(p_area).reshape(n_trials, n_windows_minus_1)
        return out

    if correction == "window":
        out = np.full_like(pvalue_matrix, np.nan, dtype=float)
        for ind_window in range(n_windows_minus_1):
            p_win = pvalue_matrix[:, :, ind_window].reshape(-1)
            out[:, :, ind_window] = bh_fdr(p_win).reshape(n_areas, n_trials)
        return out

    if correction == "trial_type":
        out = np.full_like(pvalue_matrix, np.nan, dtype=float)
        for ind_trial in range(n_trials):
            p_trial = pvalue_matrix[:, ind_trial, :].reshape(-1)
            out[:, ind_trial, :] = bh_fdr(p_trial).reshape(n_areas, n_windows_minus_1)
        return out

    if correction == "all":
        return bh_fdr(pvalue_matrix)

    return pvalue_matrix.copy()


def build_pvalue_table(
    pvalues_corrected: np.ndarray,
    *,
    trialname_list: Sequence[str],
) -> dict[str, np.ndarray]:
    """Build a trial-name keyed p-value lookup table for figure notebooks."""
    out: dict[str, np.ndarray] = {}
    for ind_trialtype, trial_name_raw in enumerate(trialname_list):
        trial_name = trial_name_raw.replace("-", "_").replace(" ", "_")
        out[trial_name] = np.squeeze(pvalues_corrected[:, ind_trialtype, :])
    return out


def plot_paired_session_summary(
    *,
    performance_by_area: Mapping[str, np.ndarray],
    pvalues_corrected: np.ndarray,
    region_list: Sequence[str],
    plot_labels: Sequence[str],
    ylabel: str,
    y_limits: tuple[float, float],
    y_ticks: Sequence[float],
    figure_size: tuple[float, float] = (15.8, 3.6),
    dpi: int = 170,
    mean_color: Any = "#1f4e79",
    session_color: Any = (0.72, 0.72, 0.72),
    xtick_rotation: float = 0.0,
    xtick_ha: str = "center",
    significance_mode: str = "simple",
    significance_kwargs: Mapping[str, Any] | None = None,
) -> tuple[Any, np.ndarray]:
    """Plot the paired session summaries used by the optogenetic figures.

    Parameters
    ----------
    performance_by_area:
        Mapping from area name to `(conditions, sessions)` summary matrix.
    pvalues_corrected:
        Corrected p-value tensor indexed by area, trial type, and comparison
        window.
    region_list:
        Areas to render from left to right.
    plot_labels:
        Condition labels shown on the x-axis.
    ylabel:
        Shared y-axis label.
    y_limits, y_ticks:
        Shared axis-range settings for all subplots.
    figure_size, dpi:
        Figure sizing controls.
    mean_color, session_color:
        Styling for the mean trace and the individual session traces.
    xtick_rotation, xtick_ha:
        Tick-label styling controls.
    significance_mode:
        Whether to draw one shared-row bracket set or stacked brackets.
    significance_kwargs:
        Additional keyword arguments forwarded to the significance helper.

    Returns
    -------
    tuple[Any, np.ndarray]
        Figure handle and one-dimensional axes array.
    """
    fig, axs = plt.subplots(1, len(region_list), figsize=figure_size, dpi=dpi, constrained_layout=False, facecolor="white")
    axs = np.asarray(axs).reshape(-1)
    x = np.arange(1, len(plot_labels) + 1)
    sig_kwargs = dict(significance_kwargs or {})

    for ind_area, current_area in enumerate(region_list):
        ax = axs[ind_area]
        performance = np.asarray(performance_by_area[current_area], dtype=float)
        n_sessions = performance.shape[1]

        # Draw individual sessions first so the group summary remains legible
        # on top of the paired trajectories.
        for sess_idx in range(n_sessions):
            ax.plot(
                x,
                performance[:, sess_idx],
                "o-",
                color=session_color,
                linewidth=0.9,
                markersize=2.8,
                markerfacecolor="white",
                markeredgewidth=0.6,
                alpha=0.9,
                zorder=1,
            )

        mu = nanmean_no_warn(performance, axis=1)
        sd = nanstd_no_warn(performance, axis=1)
        ax.errorbar(
            x,
            mu,
            yerr=sd,
            fmt="-o",
            color=mean_color,
            markersize=6.5,
            linewidth=2.2,
            capsize=2.5,
            capthick=1.1,
            markerfacecolor="white",
            markeredgecolor=mean_color,
            markeredgewidth=1.2,
            zorder=3,
        )

        ax.set_ylim(*y_limits)
        ax.set_xlim(0.5, len(plot_labels) + 0.5)
        ax.set_xticks(x)
        ax.set_xticklabels(plot_labels, rotation=xtick_rotation, ha=xtick_ha)
        ax.set_yticks(list(y_ticks))
        ax.set_title(f"{current_area}\n$n$={n_sessions}", fontsize=11.5, pad=8)
        if ind_area == 0:
            ax.set_ylabel(ylabel, fontsize=11)
        else:
            ax.set_yticklabels([])

        if significance_mode == "stacked":
            add_significance_stacked(ax, np.squeeze(pvalues_corrected[ind_area, 0, :]), [2, 3, 4], **sig_kwargs)
        else:
            add_significance_simple(ax, np.squeeze(pvalues_corrected[ind_area, 0, :]), [2, 3, 4], **sig_kwargs)

        style_plot_axis(
            ax,
            grid_axis="y",
            grid_color="0.92",
            grid_linewidth=0.8,
            grid_alpha=1.0,
            spine_width=1.2,
            tick_width=1.0,
            tick_length=2.5,
            tick_labelsize=9.5,
            hide_top=True,
            hide_right=True,
        )
        ax.spines["left"].set_color("0.15")
        ax.spines["bottom"].set_color("0.15")
        ax.tick_params(color="0.15")
        ax.set_axisbelow(True)

    return fig, axs