"""Public helper module.

Conceptually, this module keeps the lightweight data-normalization helpers used
directly by structured figure notebooks.

It exists as a separate unit so small figure-facing utilities remain easy to
import while a few long-standing coercion helpers stay available on stable
paths.

It connects notebook plotting code to small masking, summary-statistic, and
annotation helpers without pulling in heavier analysis modules.
"""

from __future__ import annotations

from typing import Any, Mapping, Optional

import numpy as np

from .core import clean_str, normalize_entries, to_1d as as_1d, to_1d_bool as as_bool, to_1d_str, to_dict_like
from .loading_base import align_to_len

__all__ = [
    "align_to_len",
    "as_1d",
    "as_bool",
    "clean_str",
    "completion_mask",
    "fiber_matches",
    "hex_to_rgb01",
    "mean_sem_over_axis",
    "nanmean_no_warn",
    "nanstd_no_warn",
    "normalize_entries",
    "p_to_stars",
    "to_dict_like",
]


def completion_mask(session: Mapping[str, Any], mode: str, n_trials: Optional[int] = None) -> np.ndarray:
    """Build a figure-friendly completion mask from a session mapping.

    Parameters
    ----------
    session:
        Session mapping containing at least `early_lick`, and for
        `early_licks` mode also `lick_time` and `start_time`.
    mode:
        Completion-state selector.
    n_trials:
        Optional expected trial count for shape validation.

    Returns
    -------
    np.ndarray
        Boolean trial mask with shape `(n_trials,)`.
    """
    early = as_1d(session["early_lick"], float)
    if n_trials is None:
        n_trials = int(early.size)
    if early.size != int(n_trials):
        raise ValueError(f"early_lick size {early.size} != n_trials {n_trials}")

    if mode == "completed_trials":
        return early == 0
    if mode == "early_licks":
        lick_time = as_1d(session["lick_time"], float)
        start_time = as_1d(session["start_time"], float)
        if lick_time.size != int(n_trials) or start_time.size != int(n_trials):
            raise ValueError("lick_time/start_time size mismatch")
        # Keep only true post-start early licks so missing timestamps do not
        # get promoted into this selection.
        return (0.0 < (lick_time - start_time)) & (early != 0)
    if mode == "all_trials":
        return np.ones(int(n_trials), dtype=bool)
    raise ValueError("completion_state must be one of {'completed_trials', 'early_licks', 'all_trials'}")


def fiber_matches(raw_fiber: Any, area_name: str) -> bool:
    """Return whether a fiber-location field contains a target area."""
    vals = to_1d_str(raw_fiber)
    return bool(np.any(vals == str(area_name)))


def hex_to_rgb01(hex_color: str) -> tuple[float, float, float]:
    """Convert '#RRGGBB' to matplotlib RGB tuple in [0, 1]."""
    color = hex_color.strip().lstrip("#")
    if len(color) != 6:
        raise ValueError(f"Invalid hex color: {hex_color}")
    return tuple(int(color[i : i + 2], 16) / 255.0 for i in (0, 2, 4))


def nanmean_no_warn(x: Any, axis: Optional[int] = None) -> np.ndarray:
    """np.nanmean with warnings suppressed for all-NaN slices."""
    with np.errstate(invalid="ignore", divide="ignore"):
        return np.nanmean(np.asarray(x, dtype=float), axis=axis)


def nanstd_no_warn(x: Any, axis: Optional[int] = None, ddof: int = 0) -> np.ndarray:
    """np.nanstd with warnings suppressed for all-NaN slices."""
    with np.errstate(invalid="ignore", divide="ignore"):
        return np.nanstd(np.asarray(x, dtype=float), axis=axis, ddof=ddof)


def mean_sem_over_axis(x: Any, axis: int = 0) -> tuple[np.ndarray, np.ndarray]:
    """Return mean and SEM with NaN-safe reductions.

    Code_M computes SEM with ``nanstd(X,[]) = std(X,0)`` (flag 0 = sample std,
    N-1), so use ``ddof=1`` to match the figure error bands.
    """
    arr = np.asarray(x, dtype=float)
    mean = nanmean_no_warn(arr, axis=axis)
    n = np.sum(np.isfinite(arr), axis=axis)
    std = nanstd_no_warn(arr, axis=axis, ddof=1)
    sem = std / np.sqrt(np.maximum(n, 1))
    return mean, sem


def p_to_stars(p: float) -> str:
    """Convert a p-value into the standard star annotation used in figures."""
    if not np.isfinite(p):
        return "n.s."
    if p < 1e-5:
        return "****"
    if p < 1.99e-3:
        return "***"
    if p < 1e-2:
        return "**"
    if p < 5e-2:
        return "*"
    return "n.s."