"""Public helper module.

Conceptually, this module groups the statistical tests and significance-label
rendering used across structured figure notebooks.

It exists as a separate unit so p-value calculation, multiple-comparison
helpers, and bracket drawing logic remain consistent across figures.

It connects per-panel summary values to the significance annotations rendered
on top of notebook plots.
"""

from __future__ import annotations

from typing import Any, Iterable, Sequence

import numpy as np
from scipy.stats import mannwhitneyu, norm, rankdata, t as student_t, wilcoxon

from .figure_data import p_to_stars


def signrank_p_value(no_light: Any, light: Any) -> float:
    """Return a paired Wilcoxon signed-rank p-value after NaN filtering."""
    x = np.asarray(no_light, dtype=float).reshape(-1)
    y = np.asarray(light, dtype=float).reshape(-1)
    n = min(x.size, y.size)
    x = x[:n]
    y = y[:n]
    valid = (~np.isnan(x)) & (~np.isnan(y))
    x = x[valid]
    y = y[valid]
    if x.size == 0:
        return np.nan
    # MATLAB `signrank` uses the exact distribution when the number of nonzero
    # paired differences is <= 15 and switches to the normal approximation above
    # that. scipy's method="auto" uses a different threshold (25), so choose the
    # method explicitly to match Code_M.
    n_nonzero = int(np.count_nonzero((y - x) != 0.0))
    method = "exact" if n_nonzero <= 15 else "approx"
    try:
        return float(wilcoxon(y, x, zero_method="wilcox", alternative="two-sided", method=method).pvalue)
    except ValueError:
        return 1.0


def signrank_p_value_from_diff(diff: Any) -> float:
    """Return a two-sided signed-rank p-value from one vector of paired differences.

    Parameters
    ----------
    diff:
        One-dimensional vector of paired differences.

    Returns
    -------
    float
        Two-sided p-value. Returns `nan` when no finite differences remain and
        `1.0` when all surviving differences are exactly zero.
    """

    x = np.asarray(diff, dtype=float).reshape(-1)
    x = x[np.isfinite(x)]
    if x.size == 0:
        return float("nan")

    # Exact zeros carry no rank information in the signed-rank test.
    x = x[x != 0.0]
    n = x.size
    if n == 0:
        return 1.0

    abs_x = np.abs(x)
    ranks = rankdata(abs_x, method="average")
    w_pos = float(np.sum(ranks[x > 0]))
    mu = n * (n + 1.0) / 4.0

    _, tie_counts = np.unique(abs_x, return_counts=True)
    tie_term = float(np.sum(tie_counts * (tie_counts + 1.0) * (2.0 * tie_counts + 1.0)))
    var = (n * (n + 1.0) * (2.0 * n + 1.0) - tie_term) / 24.0
    if var <= 0.0:
        return float("nan")

    z = (w_pos - mu - 0.5 * np.sign(w_pos - mu)) / np.sqrt(var)
    return float(2.0 * norm.sf(abs(z)))


def bh_fdr(pvals: Any) -> np.ndarray:
    """Apply Benjamini-Hochberg FDR correction while preserving input shape."""
    p = np.asarray(pvals, dtype=float).reshape(-1)
    out = np.full_like(p, np.nan)
    valid_idx = np.where(~np.isnan(p))[0]
    if valid_idx.size == 0:
        return out.reshape(np.asarray(pvals).shape)
    pv = p[valid_idx]
    m = pv.size
    order = np.argsort(pv)
    pv_sorted = pv[order]
    adj_sorted = pv_sorted * m / (np.arange(m) + 1)
    adj_sorted = np.minimum.accumulate(adj_sorted[::-1])[::-1]
    adj_sorted = np.clip(adj_sorted, 0.0, 1.0)
    inv_order = np.empty_like(order)
    inv_order[order] = np.arange(m)
    out[valid_idx] = adj_sorted[inv_order]
    return out.reshape(np.asarray(pvals).shape)


def add_significance_simple(
    ax: Any,
    pvals: Any,
    xs: Sequence[int],
    *,
    y: float = 0.8,
    tick: float = 0.05,
    text_offset: float = 0.02,
    color: str = "black",
    linewidth: float = 1.0,
    fontsize: float = 12.0,
) -> None:
    """Draw one significance bracket per comparison at a shared y-level.

    Parameters
    ----------
    ax:
        Target matplotlib axis.
    pvals:
        P-values mapped one-to-one to `xs`.
    xs:
        Right-hand x positions of comparisons drawn against group `1`.
    y, tick, text_offset:
        Vertical layout controls in data coordinates.
    color, linewidth, fontsize:
        Bracket and label styling.
    """

    for x, p in zip(xs, np.asarray(pvals, dtype=float).reshape(-1)):
        stars = p_to_stars(float(p))
        if not stars or stars == "n.s.":
            continue
        ax.plot([1, 1, x, x], [y - tick, y, y, y - tick], color=color, linewidth=linewidth, clip_on=False)
        ax.text((1 + x) / 2, y + text_offset, stars, ha="center", va="bottom", fontsize=fontsize)


def add_significance_stacked(
    ax: Any,
    pvals: Any,
    xs: Sequence[int],
    *,
    y_base: float = 30.0,
    level_step: float = 7.5,
    tick: float = 2.0,
    text_offset: float = 0.8,
    color: str = "black",
    linewidth: float = 1.15,
    fontsize: float = 11.5,
) -> None:
    """Draw stacked significance brackets from longest comparison to shortest."""
    significant: list[tuple[int, str]] = []
    for x, p in zip(xs, np.asarray(pvals, dtype=float).reshape(-1)):
        stars = p_to_stars(float(p))
        if stars and stars != "n.s.":
            significant.append((int(x), stars))

    if not significant:
        return

    # Draw longer comparisons first so shorter brackets stack above them in a
    # stable and readable order.
    significant.sort(key=lambda item: item[0], reverse=True)

    for level, (x, stars) in enumerate(significant):
        y = y_base + level * level_step
        ax.plot([1, 1, x, x], [y - tick, y, y, y - tick], color=color, linewidth=linewidth, clip_on=False)
        ax.text((1 + x) / 2, y + text_offset, stars, ha="center", va="bottom", fontsize=fontsize)


def pairwise_group_pvals(values: Any, groups: Any) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Return pairwise Mann-Whitney p-values for each unique group pair."""
    values = np.asarray(values, dtype=float).reshape(-1)
    groups = np.asarray(groups, dtype=float).reshape(-1)
    mask = np.isfinite(values) & np.isfinite(groups)
    values = values[mask]
    groups = groups[mask]

    uniq = np.unique(groups)
    c1, c2, pvals = [], [], []
    for i in range(len(uniq)):
        for j in range(i + 1, len(uniq)):
            g1 = values[groups == uniq[i]]
            g2 = values[groups == uniq[j]]
            if g1.size == 0 or g2.size == 0:
                p = np.nan
            else:
                try:
                    p = float(mannwhitneyu(g1, g2, alternative="two-sided").pvalue)
                except Exception:
                    p = np.nan
            c1.append(int(uniq[i]))
            c2.append(int(uniq[j]))
            pvals.append(p)
    return np.asarray(c1), np.asarray(c2), np.asarray(pvals, dtype=float)


def pairwise_kruskal_lsd_pvals(
    values: Any,
    groups: Any,
    *,
    distribution: str = "normal",
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Approximate MATLAB `kruskalwallis` + `multcompare(..., 'lsd')` on mean ranks.

    The test statistic is computed on mean-rank differences with a
    Kruskal-Wallis tie correction and converted to either a normal- or
    t-distribution p-value.
    """
    values = np.asarray(values, dtype=float).reshape(-1)
    groups = np.asarray(groups, dtype=float).reshape(-1)
    mask = np.isfinite(values) & np.isfinite(groups)
    values = values[mask]
    groups = groups[mask].astype(int)

    uniq = np.unique(groups)
    if values.size == 0 or uniq.size < 2:
        return np.asarray([], dtype=int), np.asarray([], dtype=int), np.asarray([], dtype=float)

    ranks = rankdata(values)
    n_total = float(ranks.size)
    n_by_group = np.asarray([np.sum(groups == g) for g in uniq], dtype=float)
    mean_ranks = np.asarray([np.mean(ranks[groups == g]) for g in uniq], dtype=float)

    _, tie_counts = np.unique(values, return_counts=True)
    tie_adjust = float(np.sum(tie_counts**3 - tie_counts))
    sigma_sq = n_total * (n_total + 1.0) / 12.0
    if n_total > 1.0 and tie_adjust > 0.0:
        sigma_sq -= tie_adjust / (12.0 * (n_total - 1.0))

    c1, c2, pvals = [], [], []
    df = max(int(n_total - uniq.size), 1)
    for i in range(len(uniq)):
        for j in range(i + 1, len(uniq)):
            se = np.sqrt(sigma_sq * (1.0 / n_by_group[i] + 1.0 / n_by_group[j]))
            if not np.isfinite(se) or se <= 0.0:
                p = np.nan
            else:
                stat = abs(mean_ranks[i] - mean_ranks[j]) / se
                if distribution == "t":
                    p = float(2.0 * student_t.sf(stat, df))
                elif distribution == "normal":
                    p = float(2.0 * norm.sf(stat))
                else:
                    raise ValueError(f"Unknown distribution: {distribution}")
            c1.append(int(uniq[i]))
            c2.append(int(uniq[j]))
            pvals.append(p)
    return np.asarray(c1, dtype=int), np.asarray(c2, dtype=int), np.asarray(pvals, dtype=float)


def select_pairwise_comparisons(
    x1: Any,
    x2: Any,
    pvals: Any,
    *,
    allowed_pairs: Iterable[tuple[int, int]] | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Filter pairwise comparison arrays to an allowed set of 1-based pairs."""
    xa = np.asarray(x1, dtype=int).reshape(-1)
    xb = np.asarray(x2, dtype=int).reshape(-1)
    pv = np.asarray(pvals, dtype=float).reshape(-1)
    if allowed_pairs is None:
        return xa.astype(float), xb.astype(float), pv
    allowed = {tuple(sorted((int(a), int(b)))) for a, b in allowed_pairs}
    keep = np.asarray([tuple(sorted((int(a), int(b)))) in allowed for a, b in zip(xa, xb)], dtype=bool)
    return xa[keep].astype(float), xb[keep].astype(float), pv[keep]


def add_prettify_like_pvalues(
    ax: Any,
    x1: Any,
    x2: Any,
    pvals: Any,
    *,
    text_rotation: float = 0.0,
    text_fontsize: float = 8.5,
    tick_length: float = 0.026,
    line_margin: float = 0.12,
    text_margin: float = 0.018,
    plot_non_signif: bool = False,
    nan_cutoff: float = 0.05,
    full_display_cutoff: float = 0.001,
    only_stars: bool = True,
    stars_level_1: float = 0.05,
    stars_level_2: float = 0.009,
    stars_level_3: float = 0.0005,
    y_position: float = 1.06,
    pack_levels: bool = False,
    color: str = "k",
    linewidth: float = 1.0,
    fontsize: float = 9.0,
    text_offset: float = -0.05,
) -> None:
    """Draw compact significance brackets for arbitrary pairwise comparisons.

    Parameters
    ----------
    ax:
        Target matplotlib axis.
    x1, x2:
        Left and right x positions for each comparison.
    pvals:
        P-values associated with each pair.
    text_rotation, text_fontsize:
        Text style controls.
    tick_length, line_margin, text_margin:
        Relative vertical layout controls expressed in fractions of the current
        y-range.
    plot_non_signif:
        Whether to keep and label non-significant comparisons.
    nan_cutoff, full_display_cutoff:
        Thresholds controlling when comparisons are hidden or shown as numeric
        p-values.
    only_stars:
        Whether to annotate only significance stars instead of explicit
        p-values.
    stars_level_1, stars_level_2, stars_level_3:
        P-value thresholds for one-, two-, and three-star labels.
    y_position:
        Optional explicit starting y-position, either scalar or one value per
        comparison.
    pack_levels:
        Whether non-overlapping comparisons may share one row.
    color, linewidth, fontsize, text_offset:
        Styling and text-placement controls.
    """
    x1 = np.asarray(x1, dtype=float)
    x2 = np.asarray(x2, dtype=float)
    pvals = np.asarray(pvals, dtype=float)

    if not plot_non_signif:
        mask = np.isfinite(pvals) & (pvals < nan_cutoff)
        x1 = x1[mask]
        x2 = x2[mask]
        pvals = pvals[mask]

    if pvals.size == 0:
        return

    widths = np.abs(x2 - x1)
    order = np.argsort(widths)[::-1]
    x1 = x1[order]
    x2 = x2[order]
    pvals = pvals[order]
    span_min = np.minimum(x1, x2)
    span_max = np.maximum(x1, x2)

    y0, y1 = ax.get_ylim()
    yr = float(y1 - y0)
    if yr <= 0 or not np.isfinite(yr):
        yr = 1.0

    def _bar_tops() -> tuple[np.ndarray, np.ndarray]:
        bar_x: list[float] = []
        bar_y: list[float] = []
        for container in getattr(ax, "containers", []):
            patches = getattr(container, "patches", None)
            if not patches:
                continue
            for patch in patches:
                try:
                    x = float(patch.get_x() + patch.get_width() / 2.0)
                    y = float(patch.get_y() + patch.get_height())
                except Exception:
                    continue
                if np.isfinite(x) and np.isfinite(y):
                    bar_x.append(x)
                    bar_y.append(y)
        return np.asarray(bar_x, dtype=float), np.asarray(bar_y, dtype=float)

    bar_x, bar_y = _bar_tops()

    def _get_bar_top(x: float) -> float:
        if bar_x.size == 0:
            return np.nan
        idx = np.abs(bar_x - x) < 1e-6
        if np.any(idx):
            return float(np.nanmax(bar_y[idx]))
        return np.nan

    if y_position is not None:
        y_pos_arr = np.asarray(y_position, dtype=float).reshape(-1)
        if y_pos_arr.size == 1:
            y_base = np.full(pvals.shape, float(y_pos_arr[0]), dtype=float)
        elif y_pos_arr.size == pvals.size:
            y_base = y_pos_arr.copy()
        else:
            raise ValueError("y_position must be scalar or match the number of p-values")
    else:
        y_base = np.zeros_like(pvals, dtype=float)
        # Without an explicit y-position, lift each comparison above the bars
        # it spans so dense bar plots remain readable.
        for i in range(pvals.size):
            involved_x = np.linspace(x1[i], x2[i], 10)
            bar_tops = np.asarray([_get_bar_top(x) for x in involved_x], dtype=float)
            bar_tops = bar_tops[np.isfinite(bar_tops)]
            base_y = float(np.nanmax(bar_tops)) if bar_tops.size else float(y1)
            y_base[i] = base_y + line_margin * yr

    if pack_levels:
        # Pack non-overlapping comparisons onto the same row so dense pairwise
        # panels do not grow vertically one bracket at a time.
        level_members: list[list[tuple[float, float]]] = []
        level_index = np.full(pvals.shape, -1, dtype=int)
        for i, (xa, xb) in enumerate(zip(span_min, span_max)):
            for level, intervals in enumerate(level_members):
                has_overlap = any(not (xb < ia or xa > ib) for ia, ib in intervals)
                if not has_overlap:
                    intervals.append((xa, xb))
                    level_index[i] = level
                    break
            if level_index[i] < 0:
                level_members.append([(xa, xb)])
                level_index[i] = len(level_members) - 1

        y_starts = np.zeros_like(pvals, dtype=float)
        prev_level_y = -np.inf
        for level, _ in enumerate(level_members):
            level_mask = level_index == level
            level_base = float(np.nanmax(y_base[level_mask])) if np.any(level_mask) else float(y1 + line_margin * yr)
            if level == 0:
                level_y = level_base
            else:
                level_y = max(level_base, prev_level_y + line_margin * yr)
            y_starts[level_mask] = level_y
            prev_level_y = level_y
    else:
        y_starts = np.zeros_like(pvals, dtype=float)
        for i in range(pvals.size):
            if i == 0:
                y_starts[i] = float(y_base[i])
            else:
                y_starts[i] = max(float(y_base[i]), y_starts[i - 1] + line_margin * yr)

    def _format_pvalue_text(p: float) -> str:
        if not np.isfinite(p) or p >= nan_cutoff:
            return "n.s."
        if only_stars:
            if p < stars_level_3:
                return "***"
            if p < stars_level_2:
                return "**"
            if p < stars_level_1:
                return "*"
            return ""
        if p < full_display_cutoff:
            return f"p < {full_display_cutoff:g}"
        return f"p = {p:.3f}"

    for i, (p, xa, xb) in enumerate(zip(pvals, x1, x2)):
        y_curr = float(y_starts[i])
        txt = _format_pvalue_text(float(p))
        ax.plot([xa, xb], [y_curr, y_curr], color=color, linewidth=linewidth, clip_on=False)
        ax.plot([xa, xa], [y_curr, y_curr - tick_length * yr], color=color, linewidth=linewidth, clip_on=False)
        ax.plot([xb, xb], [y_curr, y_curr - tick_length * yr], color=color, linewidth=linewidth, clip_on=False)
        ax.text(
            (xa + xb) / 2.0,
            y_curr + text_offset,
            txt,
            ha="center",
            va="bottom",
            rotation=text_rotation,
            fontsize=text_fontsize if text_fontsize is not None else fontsize,
        )

    new_ymax = max(float(ax.get_ylim()[1]), float(np.nanmax(y_starts)) + 2.0 * text_margin * yr)
    ax.set_ylim(float(ax.get_ylim()[0]), new_ymax)