"""Appendix A1 — LightGCN hyperparameter sweep (dim x layers) F1 heatmap."""
from pathlib import Path
import re
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from style import apply, save, PALETTE as C, COL2  # noqa: E402

KEY = "figA1_hyperparam"
TITLE = "Appendix Figure A1. LightGCN hyperparameter sweep"


def make(root, out):
    apply()
    csv = root / "validation_runs" / "dynamic_summary.csv"
    df = pd.read_csv(csv)
    df = df[df.split == "dynamic_seed202"].copy()

    def parse(run):
        m = re.search(r"l(\d)d(\d+)", run)
        return pd.Series([int(m.group(1)), int(m.group(2))] if m else [np.nan, np.nan])

    df[["layers", "dim"]] = df.run.apply(parse)
    df = df.dropna(subset=["layers"])
    piv = df.groupby(["dim", "layers"]).f1.max().unstack()

    fig, ax = plt.subplots(figsize=(COL2 * 0.62, 3.4))
    sns.heatmap(piv, annot=True, fmt=".5f", cmap="viridis", linewidths=0.5, linecolor="white",
                cbar_kws={"label": "best validation F1"}, ax=ax)
    try:
        col = list(piv.columns).index(2)
        row = list(piv.index).index(512)
        ax.add_patch(plt.Rectangle((col, row), 1, 1, fill=False, edgecolor=C[3], lw=2.5))
    except ValueError:
        pass
    ax.set_xlabel("propagation layers"); ax.set_ylabel("embedding dim")
    ax.set_title("LightGCN sweep — chosen config: 2 layers, dim 512", fontsize=9)
    save(fig, KEY, out)
    return dict(key=KEY, title=TITLE, status="ok", files=[f"{KEY}.pdf", f"{KEY}.png", f"{KEY}.svg"],
                sources=[str(csv)],
                caption=(
        "LightGCN hyperparameter sweep (best validation F1 per dim×layers cell, seed=202). "
        "Two propagation layers with a 512-d embedding is optimal (0.93858); a single layer "
        "under-fits and very large embeddings over-fit. This is the configuration used throughout."))


if __name__ == "__main__":
    from style import ensure_dirs
    r = make(Path("."), ensure_dirs(Path(".")))
    print(r["key"], r["status"])