"""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 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")