src/dropout_decay/experiments/reports.py

"""
Derived from Andrej Karpathy's nanochat project.

MIT License

Copyright (c) 2025 Andrej Karpathy

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
"""

from __future__ import annotations

from collections import defaultdict
from pathlib import Path
import statistics


def write_screen_markdown_summary(output_dir: Path, rows: list[dict]) -> None:
    if not rows:
        return
    static_rows = [
        row
        for row in rows
        if row["run_mode"] == "screen_static" and row["condition_kind"] == "static"
    ]
    if not static_rows:
        return

    by_model_prefix_rate: dict[tuple[str, int, float], list[dict]] = defaultdict(list)
    for row in rows:
        if row["run_mode"] == "screen_static" and row["condition_kind"] == "static":
            by_model_prefix_rate[
                (
                    row["model_name"],
                    int(row["token_limit"]),
                    float(row["dropout_initial"]),
                )
            ].append(row)

    aggregates: list[dict] = []
    for (model_name, prefix, dropout), group_rows in by_model_prefix_rate.items():
        first = group_rows[0]
        val_losses = [float(row["val_eval_loss"]) for row in group_rows]
        train_losses = [float(row["train_eval_loss"]) for row in group_rows]
        gaps = [float(row["generalization_gap"]) for row in group_rows]
        aggregates.append(
            {
                "model_name": model_name,
                "token_limit": prefix,
                "dropout_initial": dropout,
                "n": len(group_rows),
                "mean_val_eval_loss": statistics.fmean(val_losses),
                "std_val_eval_loss": statistics.stdev(val_losses)
                if len(val_losses) > 1
                else 0.0,
                "mean_train_eval_loss": statistics.fmean(train_losses),
                "std_train_eval_loss": statistics.stdev(train_losses)
                if len(train_losses) > 1
                else 0.0,
                "mean_generalization_gap": statistics.fmean(gaps),
                "std_generalization_gap": statistics.stdev(gaps)
                if len(gaps) > 1
                else 0.0,
                "parameters": int(first["parameters"]),
                "n_layer": int(first["n_layer"]),
                "n_head": int(first["n_head"]),
                "n_embd": int(first["n_embd"]),
                "block_size": int(first["model_config"]["block_size"]),
                "vocab_size": int(first["model_config"]["vocab_size"]),
                "tokens_seen": int(first["tokens_seen"]),
                "seeds": sorted({int(row["seed"]) for row in group_rows}),
            }
        )

    by_model: dict[str, list[dict]] = defaultdict(list)
    for row in aggregates:
        by_model[row["model_name"]].append(row)

    model_rows = []
    for model_name, model_group in by_model.items():
        first = model_group[0]
        seeds = sorted({seed for row in model_group for seed in row["seeds"]})
        model_rows.append(
            {
                "model_name": model_name,
                "parameters": first["parameters"],
                "n_layer": first["n_layer"],
                "n_head": first["n_head"],
                "n_embd": first["n_embd"],
                "block_size": first["block_size"],
                "vocab_size": first["vocab_size"],
                "seeds": seeds,
            }
        )

    lines = [
        "# Static Dropout Screen Summary",
        "",
        f"Run directory: `{output_dir}`",
        "",
        "## Models",
        "",
        "| Model | Params | Layers | Heads | Embedding | Block | Vocab | Seeds |",
        "|---|---:|---:|---:|---:|---:|---:|---|",
    ]
    for model in sorted(model_rows, key=lambda item: item["parameters"]):
        lines.append(
            "| "
            f"`{model['model_name']}` | {model['parameters']:,} | "
            f"{model['n_layer']} | {model['n_head']} | {model['n_embd']} | "
            f"{model['block_size']} | {model['vocab_size']} | "
            f"{', '.join(str(seed) for seed in model['seeds'])} |"
        )
    lines.extend(
        [
            "",
            "## Best Dropout By Model And Prefix",
            "",
            "| Model | Prefix tokens | Effective epochs | Best dropout | Mean val loss | Val std | Mean train loss | Mean gap | Plateau/bracket note |",
            "|---|---:|---:|---:|---:|---:|---:|---:|---|",
        ]
    )
    for model_name, model_group in sorted(by_model.items()):
        by_prefix: dict[int, list[dict]] = defaultdict(list)
        for row in model_group:
            by_prefix[int(row["token_limit"])].append(row)
        for prefix, prefix_rows in sorted(by_prefix.items()):
            best = min(prefix_rows, key=lambda row: row["mean_val_eval_loss"])
            rates = [float(row["dropout_initial"]) for row in prefix_rows]
            eff_epochs = float(best["tokens_seen"]) / prefix
            if best["dropout_initial"] == max(rates):
                note = "not bracketed; best at top of tested grid"
            elif best["dropout_initial"] == min(rates):
                note = "not bracketed; best at bottom of tested grid"
            else:
                note = "bracketed by tested grid"
            lines.append(
                "| "
                f"`{model_name}` | {prefix:,} | {eff_epochs:.2f} | "
                f"{best['dropout_initial']:.2f} | "
                f"{best['mean_val_eval_loss']:.4f} | "
                f"{best['std_val_eval_loss']:.4f} | "
                f"{best['mean_train_eval_loss']:.4f} | "
                f"{best['mean_generalization_gap']:.4f} | {note} |"
            )

    for model_name, model_group in sorted(by_model.items()):
        by_prefix = defaultdict(list)
        for row in model_group:
            by_prefix[int(row["token_limit"])].append(row)
        lines.extend(
            [
                "",
                f"## Model `{model_name}`",
            ]
        )
        for prefix, prefix_rows in sorted(by_prefix.items()):
            eff_epochs = float(prefix_rows[0]["tokens_seen"]) / prefix
            lines.extend(
                [
                    "",
                    f"### Prefix {prefix:,} Tokens ({eff_epochs:.2f} Effective Epochs)",
                    "",
                    "| Dropout | N | Mean val loss | Val std | Mean train loss | Train std | Mean gap | Gap std | Sampled tokens | Params |",
                    "|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|",
                ]
            )
            for row in sorted(prefix_rows, key=lambda item: item["dropout_initial"]):
                lines.append(
                    "| "
                    f"{row['dropout_initial']:.2f} | {row['n']} | "
                    f"{row['mean_val_eval_loss']:.4f} | "
                    f"{row['std_val_eval_loss']:.4f} | "
                    f"{row['mean_train_eval_loss']:.4f} | "
                    f"{row['std_train_eval_loss']:.4f} | "
                    f"{row['mean_generalization_gap']:.4f} | "
                    f"{row['std_generalization_gap']:.4f} | "
                    f"{int(row['tokens_seen']):,} | {int(row['parameters']):,} |"
                )

    output = "\n".join(lines) + "\n"
    (output_dir / "RESULT_SUMMARY.md").write_text(output, encoding="utf-8")


def svg_escape(value: object) -> str:
    return (
        str(value)
        .replace("&", "&")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
    )


def write_dropout_curve_svg(output_dir: Path, summary: list[dict]) -> None:
    rows = [
        row
        for row in summary
        if row["run_mode"] == "screen_static" and row["condition_kind"] == "static"
    ]
    if not rows:
        return

    grouped: dict[tuple[str, int], list[dict]] = defaultdict(list)
    model_params: dict[str, int] = {}
    for row in rows:
        model_name = row["model_name"]
        grouped[(model_name, int(row["token_limit"]))].append(row)
        model_params[model_name] = int(row["parameters"])

    models = sorted(model_params, key=lambda name: model_params[name])
    prefixes = sorted({int(row["token_limit"]) for row in rows})
    panel_w, panel_h = 230, 170
    margin_l, margin_t = 58, 34
    plot_w, plot_h = 142, 94
    gap_x, gap_y = 18, 38
    width = margin_l + len(prefixes) * panel_w + gap_x
    height = 70 + len(models) * (panel_h + gap_y)
    colors = ["#1f77b4", "#d62728", "#2ca02c", "#9467bd", "#ff7f0e"]

    parts = [
        f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">',
        "<style>",
        "text{font-family:Arial,Helvetica,sans-serif;fill:#111827}",
        ".small{font-size:10px}.label{font-size:11px}.title{font-size:15px;font-weight:700}",
        ".axis{stroke:#374151;stroke-width:1}.grid{stroke:#e5e7eb;stroke-width:1}.line{fill:none;stroke-width:2}",
        "</style>",
        '<rect width="100%" height="100%" fill="#ffffff"/>',
        '<text x="24" y="28" class="title">Static dropout law: validation loss vs dropout</text>',
        '<text x="24" y="48" class="label">Each panel uses its own y-scale. Points are one-seed means unless N &gt; 1.</text>',
    ]

    for col, prefix in enumerate(prefixes):
        x = margin_l + col * panel_w + plot_w / 2
        parts.append(
            f'<text x="{x:.1f}" y="70" text-anchor="middle" class="label">{prefix:,} prefix tokens</text>'
        )

    for row_idx, model_name in enumerate(models):
        row_y = 92 + row_idx * (panel_h + gap_y)
        parts.append(
            f'<text x="24" y="{row_y + 48}" class="label" transform="rotate(-90 24 {row_y + 48})">'
            f'{svg_escape(model_name)} ({model_params[model_name] / 1_000_000:.1f}M)</text>'
        )
        for col, prefix in enumerate(prefixes):
            panel_x = margin_l + col * panel_w
            panel_y = row_y
            curve = sorted(
                grouped.get((model_name, prefix), []),
                key=lambda item: float(item["dropout_initial"]),
            )
            if not curve:
                continue
            losses = [float(item["mean_val_eval_loss"]) for item in curve]
            min_loss, max_loss = min(losses), max(losses)
            pad = max(0.02, (max_loss - min_loss) * 0.08)
            y_min, y_max = min_loss - pad, max_loss + pad
            best = min(curve, key=lambda item: float(item["mean_val_eval_loss"]))

            def px(dropout: float) -> float:
                return panel_x + (dropout / 0.9) * plot_w

            def py(loss: float) -> float:
                scale = (loss - y_min) / (y_max - y_min)
                return panel_y + plot_h - scale * plot_h

            parts.extend(
                [
                    f'<line x1="{panel_x:.1f}" y1="{panel_y:.1f}" x2="{panel_x:.1f}" y2="{panel_y + plot_h:.1f}" class="axis"/>',
                    f'<line x1="{panel_x:.1f}" y1="{panel_y + plot_h:.1f}" x2="{panel_x + plot_w:.1f}" y2="{panel_y + plot_h:.1f}" class="axis"/>',
                    f'<line x1="{panel_x:.1f}" y1="{panel_y:.1f}" x2="{panel_x + plot_w:.1f}" y2="{panel_y:.1f}" class="grid"/>',
                    f'<text x="{panel_x:.1f}" y="{panel_y - 6:.1f}" class="small">{y_max:.2f}</text>',
                    f'<text x="{panel_x:.1f}" y="{panel_y + plot_h + 13:.1f}" class="small">{y_min:.2f}</text>',
                    f'<text x="{panel_x:.1f}" y="{panel_y + plot_h + 28:.1f}" class="small">0</text>',
                    f'<text x="{panel_x + plot_w:.1f}" y="{panel_y + plot_h + 28:.1f}" text-anchor="end" class="small">0.9</text>',
                ]
            )
            points = " ".join(
                f"{px(float(item['dropout_initial'])):.1f},{py(float(item['mean_val_eval_loss'])):.1f}"
                for item in curve
            )
            color = colors[row_idx % len(colors)]
            parts.append(f'<polyline points="{points}" class="line" stroke="{color}"/>')
            for item in curve:
                dropout = float(item["dropout_initial"])
                loss = float(item["mean_val_eval_loss"])
                radius = 4 if item is best else 2.7
                fill = "#111827" if item is best else "#ffffff"
                parts.append(
                    f'<circle cx="{px(dropout):.1f}" cy="{py(loss):.1f}" r="{radius}" fill="{fill}" stroke="{color}" stroke-width="1.5"/>'
                )
            parts.append(
                f'<text x="{panel_x + plot_w + 8:.1f}" y="{panel_y + 14:.1f}" class="small">'
                f'best p={float(best["dropout_initial"]):.2f}</text>'
            )
            parts.append(
                f'<text x="{panel_x + plot_w + 8:.1f}" y="{panel_y + 28:.1f}" class="small">'
                f'loss={float(best["mean_val_eval_loss"]):.3f}</text>'
            )

    parts.append("</svg>")
    (output_dir / "dropout_curves.svg").write_text("\n".join(parts), encoding="utf-8")


def write_stream_markdown_summary(output_dir: Path, rows: list[dict]) -> None:
    stream_rows = [row for row in rows if row["run_mode"] == "locked_stream"]
    if not stream_rows:
        return

    by_condition_stage: dict[tuple[str, int], list[dict]] = defaultdict(list)
    by_condition: dict[str, list[dict]] = defaultdict(list)
    for row in stream_rows:
        condition = row["condition"]
        by_condition_stage[(condition, int(row["stage"]))].append(row)
        by_condition[condition].append(row)

    first = stream_rows[0]
    seeds = sorted({int(row["seed"]) for row in stream_rows})
    conditions = sorted(
        by_condition,
        key=lambda name: (
            by_condition[name][0]["condition_kind"] != "anchor_decay",
            by_condition[name][0]["dropout_initial"],
            name,
        ),
    )
    stages = sorted({int(row["stage"]) for row in stream_rows})

    lines = [
        "# Locked Streaming Dropout Summary",
        "",
        f"Run directory: `{output_dir}`",
        "",
        (
            f"Model: `{first['model_name']}` causal Transformer, "
            f"{int(first['parameters']):,} parameters, {first['n_layer']} layers, "
            f"{first['n_head']} heads, {first['n_embd']} embedding dim."
        ),
        (
            f"Training per stage: {first['steps']:,} steps. "
            "Sampled tokens are cumulative in each stage row. "
            f"Seeds present: {', '.join(str(seed) for seed in seeds)}."
        ),
        "",
        "## Condition Ranking",
        "",
        "| Condition | Kind | Final dropout | Mean trajectory val loss | Final val loss | Final gap | Dropout path |",
        "|---|---|---:|---:|---:|---:|---|",
    ]

    ranking = []
    for condition in conditions:
        stage_items = []
        for stage in stages:
            group = by_condition_stage.get((condition, stage), [])
            if not group:
                continue
            stage_items.append(
                {
                    "stage": stage,
                    "token_limit": int(group[0]["token_limit"]),
                    "mean_val": statistics.fmean(
                        float(row["val_eval_loss"]) for row in group
                    ),
                    "mean_gap": statistics.fmean(
                        float(row["generalization_gap"]) for row in group
                    ),
                    "mean_dropout": statistics.fmean(
                        float(row["dropout_active_final"]) for row in group
                    ),
                    "kind": group[0]["condition_kind"],
                }
            )
        if not stage_items:
            continue
        final = max(stage_items, key=lambda item: item["stage"])
        ranking.append(
            {
                "condition": condition,
                "kind": stage_items[0]["kind"],
                "trajectory_val": statistics.fmean(item["mean_val"] for item in stage_items),
                "final_val": final["mean_val"],
                "final_gap": final["mean_gap"],
                "final_dropout": final["mean_dropout"],
                "dropout_path": " -> ".join(
                    f"{item['mean_dropout']:.2f}" for item in stage_items
                ),
            }
        )

    for item in sorted(ranking, key=lambda row: row["trajectory_val"]):
        lines.append(
            "| "
            f"`{item['condition']}` | {item['kind']} | "
            f"{item['final_dropout']:.2f} | {item['trajectory_val']:.4f} | "
            f"{item['final_val']:.4f} | {item['final_gap']:.4f} | "
            f"{item['dropout_path']} |"
        )

    lines.extend(["", "## Stage Trajectory", ""])
    for stage in stages:
        stage_groups = {
            condition: by_condition_stage[(condition, stage)]
            for condition in conditions
            if (condition, stage) in by_condition_stage
        }
        if not stage_groups:
            continue
        prefix = int(next(iter(stage_groups.values()))[0]["token_limit"])
        lines.extend(
            [
                f"### Stage {stage}: {prefix:,} Prefix Tokens",
                "",
                "| Condition | Dropout | Mean val loss | Mean train loss | Mean gap | N |",
                "|---|---:|---:|---:|---:|---:|",
            ]
        )
        for condition, group in sorted(
            stage_groups.items(),
            key=lambda item: statistics.fmean(
                float(row["val_eval_loss"]) for row in item[1]
            ),
        ):
            val = statistics.fmean(float(row["val_eval_loss"]) for row in group)
            train = statistics.fmean(float(row["train_eval_loss"]) for row in group)
            gap = statistics.fmean(float(row["generalization_gap"]) for row in group)
            dropout = statistics.fmean(
                float(row["dropout_active_final"]) for row in group
            )
            lines.append(
                "| "
                f"`{condition}` | {dropout:.2f} | {val:.4f} | "
                f"{train:.4f} | {gap:.4f} | {len(group)} |"
            )
        lines.append("")

    (output_dir / "RESULT_SUMMARY.md").write_text(
        "\n".join(lines).rstrip() + "\n",
        encoding="utf-8",
    )