"""Shared helpers for Layer-1 figure scripts.

Figures are a pure function of a committed results bundle (``results/EX/``:
``results.parquet`` + ``resolved_config.yaml``). No simulation is rerun here.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd
import yaml


def load_bundle(results_dir: str | Path) -> tuple[pd.DataFrame, dict]:
    """Load a results bundle: (long-form DataFrame, source experiment config)."""
    results_dir = Path(results_dir)
    df = pd.read_parquet(results_dir / "results.parquet")
    resolved = yaml.safe_load((results_dir / "resolved_config.yaml").read_text())
    return df, resolved["source_config"]


def load_seed_bundles(results_dir: str | Path) -> tuple[pd.DataFrame, dict]:
    """Load a per-seed bundle layout ``results_dir/s{seed}/results.parquet`` into one frame.

    Each sub-bundle gets a ``seed`` column from its directory name (the HPC array-job layout, one
    element per seed). A flat single-seed bundle is accepted too, tagged with its manifest seed.
    Returns (frame, source config of the first seed).
    """
    results_dir = Path(results_dir)
    subs = sorted(results_dir.glob("s[0-9]*/results.parquet"))
    if not subs:
        df, cfg = load_bundle(results_dir)
        if "seed" not in df.columns:
            df = df.assign(seed=int(cfg.get("seed", 1)))
        return df, cfg
    frames, cfg0 = [], None
    for p in subs:
        df, cfg = load_bundle(p.parent)
        cfg0 = cfg0 or cfg
        frames.append(df.assign(seed=int(p.parent.name[1:])))
    return pd.concat(frames, ignore_index=True), cfg0


def mean_ci(df: pd.DataFrame, by: str, value: str, ci: float = 0.95):
    """Return (index, mean, half-width) for a normal-approx CI of ``value`` grouped by ``by``."""
    from scipy import stats

    g = df.groupby(by)[value]
    mean = g.mean()
    sem = g.sem()
    z = stats.norm.ppf(0.5 + ci / 2.0)
    return mean.index.to_numpy(), mean.to_numpy(), (z * sem).to_numpy()


def savefig(fig, results_dir: str | Path, name: str) -> None:
    """Write a figure as both PNG (150 dpi) and PDF next to its results bundle."""
    results_dir = Path(results_dir)
    fig.savefig(results_dir / f"{name}.png", dpi=150, bbox_inches="tight")
    fig.savefig(results_dir / f"{name}.pdf", bbox_inches="tight")
    print(f"wrote {results_dir}/{name}.png and .pdf")


def letter_axes(fig, x: float = -0.1, y: float = 1.04, fontsize: float = 13) -> None:
    """Letter every data axes of ``fig`` A, B, C ... in reading order (top row first, left to right).

    Twin axes and colourbars share a frame with a lettered axes and are skipped. Call once, after
    every axes exists and before saving.
    """
    seen: list[tuple[float, float]] = []
    axes = []
    for ax in fig.axes:
        b = ax.get_position()
        key = (round(b.x0, 3), round(b.y0, 3))
        if key in seen or b.width < 0.05:            # twin axes / colourbars
            continue
        seen.append(key); axes.append(ax)
    axes.sort(key=lambda a: (-round(a.get_position().y0, 2), a.get_position().x0))
    for i, ax in enumerate(axes):
        ax.text(x, y, chr(ord("A") + i), transform=ax.transAxes, fontsize=fontsize, fontweight="bold",
                va="bottom", ha="left", clip_on=False)