"""Public helper module.

Conceptually, this module assembles the coordinate, color, and session-summary
helpers used by the structured optogenetic brainmap figures.

It exists as a separate unit so spatial payload loading, brainmap color logic,
and per-area optogenetic summaries stay reusable across related notebooks.

It connects saved optogenetic session payloads and area coordinates to the
arrays and colors consumed by brainmap plotting code.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List, Mapping, Optional, Sequence

import numpy as np

from .core import normalize_entries, to_session_id_str, unwrap_scalar_obj
from .figure_data import completion_mask, fiber_matches, hex_to_rgb01


VIEW_ROTATION_DEG = -90.0


def load_opto_entries(npz_path: Path) -> List[Dict[str, Any]]:
    """Load canonicalized optogenetic session entries from the processed NPZ."""
    with np.load(npz_path, allow_pickle=True) as data:
        if "optomat" not in data:
            raise KeyError("NPZ missing key 'optomat'")
        return normalize_entries(data["optomat"])


def load_area_top_coordinates(npz_path: Path, mat_path: Optional[Path] = None) -> Dict[str, np.ndarray]:
    """Decode area top coordinates from the canonical NPZ payload.

    Parameters
    ----------
    npz_path:
        Canonical NPZ source for the saved area-coordinate structure.
    mat_path:
        Ignored backward-compatibility parameter kept so older notebooks do not
        fail immediately after the NPZ-only transition.

    Returns
    -------
    dict[str, np.ndarray]
        Mapping from area name to its top-view coordinate triplet.
    """
    fields: Dict[str, Any] = {}

    del mat_path

    if npz_path.exists():
        with np.load(npz_path, allow_pickle=True) as data:
            if "area_top_coordinates" in data:
                obj = unwrap_scalar_obj(data["area_top_coordinates"])
                if isinstance(obj, np.ndarray) and obj.size == 1:
                    obj = unwrap_scalar_obj(obj.reshape(-1)[0])
                if isinstance(obj, dict):
                    fields = {str(k): v for k, v in obj.items()}
                elif hasattr(obj, "_fieldnames"):
                    fields = {str(f): getattr(obj, f) for f in (obj._fieldnames or [])}
                elif isinstance(obj, np.void) and obj.dtype.names:
                    fields = {str(n): obj[n] for n in obj.dtype.names}

    if not fields:
        raise ValueError("Could not decode Area_Top_Coordinates from NPZ")

    out: Dict[str, np.ndarray] = {}
    for key, value in fields.items():
        arr = np.asarray(unwrap_scalar_obj(value), dtype=float).reshape(-1)
        if arr.size >= 3:
            out[key] = arr[:3]

    if not out:
        raise ValueError("Decoded Area_Top_Coordinates is empty or malformed")
    return out


def build_brainmap_colormap(vmin: float, vmax: float) -> np.ndarray:
    """Build the custom blue-white-red colormap used by the brainmap figures."""
    n_neg = int(abs(vmin))
    n_pos = int(abs(vmax))
    white = np.array([1.0, 1.0, 1.0], dtype=float)
    blue = np.array([0.0, 0.0, 1.0], dtype=float)
    red = np.array([1.0, 0.0, 0.0], dtype=float)
    map_neg = np.vstack([np.linspace(white[i], blue[i], n_neg) for i in range(3)]).T
    map_pos = np.vstack([np.linspace(white[i], red[i], n_pos) for i in range(3)]).T
    return np.vstack([np.flipud(map_neg), map_pos])


def map_value_to_color(value: float, cmap: np.ndarray, vmin: float, vmax: float) -> np.ndarray:
    """Map a scalar value onto the custom brainmap colormap."""
    if np.isnan(value):
        return np.array([0.6, 0.6, 0.6], dtype=float)
    normalized = (value - vmin) / (vmax - vmin)
    idx = int(np.rint(normalized * (cmap.shape[0] - 1)))
    idx = max(min(idx, cmap.shape[0] - 1), 0)
    return cmap[idx, :]


def draw_brainmap_reference(ax: Any, *, marker_size: float = 15.0, line_width: float = 2.0) -> None:
    """Draw the top-view reference frame used in the optogenetic brainmaps."""
    ax.set_xlim(0, 1000)
    ax.set_ylim(0, 600)
    ax.set_zlim(0, 500)
    ax.set_axis_off()
    ax.invert_zaxis()
    ax.view_init(elev=90, azim=90 + VIEW_ROTATION_DEG)

    bregma = np.asarray([[540.0, 0.0, 570.0]], dtype=float)
    ax.plot3D(bregma[:, 0], bregma[:, 2], bregma[:, 1], "+", color="k", markersize=marker_size, linewidth=line_width)

    xlin = np.asarray([[740, 0, 570], [690, 0, 570]], dtype=float)
    ylin = np.asarray([[740, 0, 570], [740, 0, 520]], dtype=float)
    ax.plot3D(xlin[:, 0], xlin[:, 2], xlin[:, 1], "-", color="k", linewidth=line_width)
    ax.plot3D(ylin[:, 0], ylin[:, 2], ylin[:, 1], "-", color="k", linewidth=line_width)


def compute_diff_mean_for_trial_type(
    optomat: Sequence[Mapping[str, Any]],
    *,
    current_trial_type: int,
    region_list: Sequence[str],
    period_list: Sequence[str],
    completion_state: str,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Compute light-minus-nolight lick deltas for one trial type.

    Parameters
    ----------
    optomat:
        Session-level optogenetic payloads.
    current_trial_type:
        Trial-type code to summarize.
    region_list:
        Areas evaluated on the brainmap.
    period_list:
        Optogenetic windows evaluated against the `nolight` reference window.
    completion_state:
        Trial-completion mode used to filter the session trials.

    Returns
    -------
    tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
        Per-area mean differences, SEM values, number of sessions, and number
        of mice contributing to each area/window pair.
    """
    n_area = len(region_list)
    n_win = len(period_list)

    diff_mean = np.full((n_area, n_win), np.nan, dtype=float)
    diff_sem = np.full((n_area, n_win), np.nan, dtype=float)
    num_session = np.zeros(n_area, dtype=float)
    num_mice = np.full((n_area, n_win), np.nan, dtype=float)

    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)]
        if len(fiberlist) == 0:
            continue

        perf = np.full((n_win, len(fiberlist)), np.nan, dtype=float)
        mice_name = np.empty((n_win, len(fiberlist)), dtype=object)

        for sess_i, ind_session in enumerate(fiberlist):
            session = optomat[ind_session]
            required = ["trial_type", "lick_flag", "opto_window", "early_lick", "session_id"]
            missing = [k for k in required if k not in session]
            if missing:
                sid = to_session_id_str(session.get("session_id", f"idx_{ind_session}"))
                raise KeyError(f"Session {sid}: missing fields {missing}")

            trial = np.asarray(unwrap_scalar_obj(session["trial_type"]), dtype=float).reshape(-1)
            lick = np.asarray(unwrap_scalar_obj(session["lick_flag"]), dtype=float).reshape(-1) == 1
            windows = np.asarray([str(v).strip() for v in np.asarray(unwrap_scalar_obj(session["opto_window"])).reshape(-1)], dtype=object)
            n_trials = int(trial.size)

            if lick.size != n_trials or windows.size != n_trials:
                sid = to_session_id_str(session["session_id"])
                raise ValueError(f"Session {sid}: trial_type={n_trials}, lick_flag={lick.size}, opto_window={windows.size}")

            current_trialtype_ind = trial == float(current_trial_type)
            completion_state_ind = completion_mask(session, completion_state, n_trials)

            sid = to_session_id_str(session["session_id"])
            mouse = sid[4:9] if len(sid) >= 9 else sid

            # Compare each stimulation window to its no-light baseline within
            # the same session before averaging across sessions.
            for ind_window, current_window in enumerate(period_list):
                current_optocondition_ind = windows == current_window
                nolight_optocondition_ind = windows == "nolight"

                curr_trial_ind = completion_state_ind & current_trialtype_ind & current_optocondition_ind
                nolight_trial_ind = completion_state_ind & current_trialtype_ind & nolight_optocondition_ind

                denom_light = int(np.sum(curr_trial_ind))
                denom_nolight = int(np.sum(nolight_trial_ind))
                if denom_light == 0 or denom_nolight == 0:
                    plick = np.nan
                else:
                    p_light = (np.sum(curr_trial_ind & lick) / float(denom_light)) * 100.0
                    p_nolight = (np.sum(nolight_trial_ind & lick) / float(denom_nolight)) * 100.0
                    plick = p_light - p_nolight

                perf[ind_window, sess_i] = plick
                mice_name[ind_window, sess_i] = mouse

        diff_mean[ind_area, :] = np.nanmean(perf, axis=1)
        # MATLAB nanstd uses the N-1 (sample) normalization; match it with ddof=1.
        diff_sem[ind_area, :] = np.nanstd(perf, axis=1, ddof=1) / np.sqrt(float(perf.shape[1]))
        num_session[ind_area] = float(perf.shape[1])

        for w in range(n_win):
            mice_row = [m for m in mice_name[w, :].tolist() if m not in (None, "")]
            if mice_row:
                num_mice[ind_area, w] = float(len(np.unique(mice_row)))

    return diff_mean, diff_sem, num_session, num_mice


def area_hex_colors(region_list: Sequence[str], color_map: Mapping[str, str]) -> list[tuple[float, float, float]]:
    """Resolve area colors as matplotlib RGB tuples."""
    return [hex_to_rgb01(color_map.get(area, "#808080")) for area in region_list]