#!/usr/bin/env python3
"""Generate publication-quality CodePin SFT training figures from archived logs."""

from __future__ import annotations

import argparse
import csv
import json
import math
import re
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import FuncFormatter, MaxNLocator


ANSI_RE = re.compile(r"\x1b\[[0-9;]*[A-Za-z]")
STEP_BLOCK_RE = re.compile(r"Step\s+(\d+):\s*\n(.*?\{.*?\})", re.DOTALL)
PAIR_RE = re.compile(r"'([^']+)':\s*'?([-+0-9.eE]+)'?")

BLUE = "#0072B2"
SKY = "#56B4E9"
GREEN = "#009E73"
ORANGE = "#E69F00"
VERMILLION = "#D55E00"
PINK = "#CC79A7"
GRAY = "#7A7A7A"
LIGHT_GRAY = "#D6D9DC"
INK = "#24292F"


def parse_log(path: Path) -> dict[int, dict[str, float]]:
    text = ANSI_RE.sub("", path.read_text(encoding="utf-8", errors="replace"))
    records: dict[int, dict[str, float]] = {}
    for _, block in STEP_BLOCK_RE.findall(text):
        values = {key: float(value) for key, value in PAIR_RE.findall(block)}
        if "train/global_step" in values:
            records[int(values["train/global_step"])] = values
    return records


def rolling(values: np.ndarray, window: int, reducer=np.median) -> np.ndarray:
    output = np.empty_like(values, dtype=float)
    radius = window // 2
    for index in range(len(values)):
        lo = max(0, index - radius)
        hi = min(len(values), index + radius + 1)
        output[index] = reducer(values[lo:hi])
    return output


def ema(values: np.ndarray, alpha: float = 0.08) -> np.ndarray:
    output = np.empty_like(values, dtype=float)
    output[0] = values[0]
    for index in range(1, len(values)):
        output[index] = alpha * values[index] + (1.0 - alpha) * output[index - 1]
    return output


def scheduled_lr(steps: np.ndarray, peak: float = 5e-5, warmup: int = 71, total: int = 710) -> np.ndarray:
    result = np.zeros_like(steps, dtype=float)
    for index, step in enumerate(steps):
        if step <= warmup:
            result[index] = peak * step / warmup
        else:
            progress = min(1.0, (step - warmup) / (total - warmup))
            result[index] = peak * 0.5 * (1.0 + math.cos(math.pi * progress))
    return result


def configure_style() -> None:
    plt.rcParams.update(
        {
            "font.family": "sans-serif",
            "font.sans-serif": ["Inter", "Arial", "DejaVu Sans"],
            "font.size": 8.8,
            "axes.titlesize": 10,
            "axes.titleweight": "bold",
            "axes.labelsize": 8.5,
            "legend.fontsize": 7.4,
            "legend.frameon": False,
            "figure.dpi": 180,
            "savefig.dpi": 300,
            "savefig.bbox": "tight",
            "axes.spines.top": False,
            "axes.spines.right": False,
            "axes.edgecolor": "#555555",
            "axes.linewidth": 0.7,
            "axes.grid": True,
            "grid.alpha": 0.18,
            "grid.linestyle": "-",
            "grid.linewidth": 0.55,
            "lines.linewidth": 1.6,
            "xtick.labelsize": 7.8,
            "ytick.labelsize": 7.8,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )


def panel_label(ax: plt.Axes, label: str) -> None:
    ax.text(
        0.015,
        0.975,
        label,
        transform=ax.transAxes,
        fontweight="bold",
        fontsize=9.5,
        va="top",
        ha="left",
        bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.78, "pad": 1.2},
        zorder=10,
    )


def save_metrics_csv(path: Path, steps: np.ndarray, arrays: dict[str, np.ndarray]) -> None:
    fieldnames = ["step", *arrays.keys()]
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for index, step in enumerate(steps):
            row = {"step": int(step)}
            for key, values in arrays.items():
                value = values[index]
                row[key] = str(value) if isinstance(value, str) else float(value)
            writer.writerow(row)


def stats(values: np.ndarray) -> dict[str, float]:
    return {
        "min": float(np.min(values)),
        "median": float(np.median(values)),
        "p95": float(np.percentile(values, 95)),
        "max": float(np.max(values)),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, default=Path("figures"))
    args = parser.parse_args()
    run_dir = args.run_dir.resolve()
    output_dir = args.output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    initial = parse_log(run_dir / "logs/main.log")
    successful_resume = parse_log(run_dir / "logs/main_resume3.log")
    records = {step: values for step, values in initial.items() if step <= 400}
    records.update({step: values for step, values in successful_resume.items() if step >= 401})
    expected = set(range(1, 711))
    if set(records) != expected:
        missing = sorted(expected - set(records))
        extra = sorted(set(records) - expected)
        raise RuntimeError(f"Expected exactly steps 1..710; missing={missing}, extra={extra}")

    steps = np.arange(1, 711)

    def series(key: str) -> np.ndarray:
        return np.asarray([records[int(step)][key] for step in steps], dtype=float)

    loss = series("train/loss")
    grad_norm = series("train/grad_norm")
    padded_length = series("train/batch_padded_seq_len")
    actual_tokens = series("train/actual_num_tokens")
    throughput = series("train/tokens_per_second_per_gpu")
    step_time = series("timing/step")
    forward_backward_time = series("timing/forward_backward")
    optimizer_time = series("timing/optim_step")
    data_time = series("timing/data_loading")
    cumulative_tokens = np.cumsum(actual_tokens)
    padding_efficiency = actual_tokens / (8.0 * padded_length)
    learning_rate = scheduled_lr(steps)

    arrays = {
        "loss": loss,
        "grad_norm": grad_norm,
        "padded_sequence_length": padded_length,
        "actual_tokens": actual_tokens,
        "padding_efficiency": padding_efficiency,
        "tokens_per_second_per_gpu": throughput,
        "step_time_seconds": step_time,
        "forward_backward_seconds": forward_backward_time,
        "optimizer_step_seconds": optimizer_time,
        "data_loading_seconds": data_time,
        "cumulative_actual_tokens": cumulative_tokens,
        "reconstructed_learning_rate": learning_rate,
    }
    save_metrics_csv(output_dir / "sft_training_metrics.csv", steps, arrays)

    long_mask = padded_length >= 7500
    step_507_index = 506
    step_640_index = 639
    summary = {
        "status": "PASS",
        "trajectory_construction": {
            "description": "one continuous sequence of unique successful optimizer steps",
            "unique_steps": 710,
            "duplicate_failed_attempts_excluded": True,
        },
        "optimization": {
            "loss": stats(loss),
            "loss_first_50_median": float(np.median(loss[:50])),
            "loss_last_50_median": float(np.median(loss[-50:])),
            "loss_final": float(loss[-1]),
            "final_eval_loss": 0.2157,
            "grad_norm": stats(grad_norm),
            "grad_norm_last_50_median": float(np.median(grad_norm[-50:])),
        },
        "data_and_context": {
            "total_actual_tokens": int(cumulative_tokens[-1]),
            "padded_sequence_length": stats(padded_length),
            "batches_at_least_7500_tokens": int(long_mask.sum()),
            "max_length_step": int(steps[np.argmax(padded_length)]),
            "padding_efficiency": stats(padding_efficiency),
        },
        "systems": {
            "throughput_tokens_per_second_per_gpu": stats(throughput),
            "step_time_seconds": stats(step_time),
            "optimizer_time_seconds": stats(optimizer_time),
            "sequence_length_vs_step_time_pearson_r": float(np.corrcoef(padded_length, step_time)[0, 1]),
            "sequence_length_vs_throughput_pearson_r": float(np.corrcoef(padded_length, throughput)[0, 1]),
        },
        "critical_batches": {
            "step_507": {key: float(values[step_507_index]) for key, values in {
                "loss": loss,
                "grad_norm": grad_norm,
                "padded_sequence_length": padded_length,
                "actual_tokens": actual_tokens,
                "tokens_per_second_per_gpu": throughput,
                "step_time_seconds": step_time,
            }.items()},
            "step_640_max_length": {key: float(values[step_640_index]) for key, values in {
                "loss": loss,
                "grad_norm": grad_norm,
                "padded_sequence_length": padded_length,
                "tokens_per_second_per_gpu": throughput,
                "step_time_seconds": step_time,
            }.items()},
        },
        "limitations": [
            "No continuous validation series was logged because eval_interval=0; only final eval_loss=0.2157 is available.",
            "No per-step GPU memory/utilization telemetry was recorded, so memory curves are not inferred.",
            "Learning rate is reconstructed from the declared 71-step warmup and 710-step cosine schedule.",
        ],
    }
    (output_dir / "sft_training_summary.json").write_text(
        json.dumps(summary, indent=2) + "\n", encoding="utf-8"
    )

    configure_style()
    fig, axes = plt.subplots(2, 3, figsize=(10.5, 6.25), constrained_layout=False)
    fig.subplots_adjust(left=0.075, right=0.96, bottom=0.09, top=0.82, wspace=0.36, hspace=0.42)
    axes = axes.ravel()

    ax = axes[0]
    ax.plot(steps, loss, color=SKY, alpha=0.28, linewidth=0.7, label="per-step")
    ax.plot(steps, ema(loss), color=BLUE, linewidth=1.9, label="EMA (α=0.08)")
    ax.axhline(0.2157, color=VERMILLION, linestyle=(0, (4, 2)), linewidth=1.1, label="final eval = 0.2157")
    ax.scatter([710], [loss[-1]], color=BLUE, edgecolor="white", linewidth=0.6, s=28, zorder=5)
    ax.annotate("0.1577", (710, loss[-1]), xytext=(-7, 7), textcoords="offset points", ha="right", fontsize=7, color=BLUE)
    ax.set_title("Optimization objective")
    ax.set_ylabel("Token-level SFT loss")
    ax.set_ylim(0, max(0.62, np.percentile(loss, 99) * 1.08))
    ax.legend(loc="upper right", ncol=1)
    panel_label(ax, "a")

    ax = axes[1]
    ax.plot(steps, grad_norm, color=ORANGE, alpha=0.25, linewidth=0.7, label="per-step")
    ax.plot(steps, rolling(grad_norm, 25), color=VERMILLION, linewidth=1.8, label="rolling median (25)")
    peak_index = int(np.argmax(grad_norm))
    ax.scatter(steps[peak_index], grad_norm[peak_index], color=VERMILLION, s=22, zorder=5)
    ax.annotate(
        f"peak {grad_norm[peak_index]:.1f}",
        (steps[peak_index], grad_norm[peak_index]),
        xytext=(14, -2),
        textcoords="offset points",
        fontsize=7,
        color=VERMILLION,
    )
    ax.set_yscale("log")
    ax.set_title("Gradient stability")
    ax.set_ylabel("Global gradient norm (log)")
    ax.legend(loc="upper right")
    panel_label(ax, "b")

    ax = axes[2]
    ax.axhspan(7500, 8192, color=VERMILLION, alpha=0.08, linewidth=0)
    ax.plot(steps, padded_length, color=BLUE, alpha=0.62, linewidth=0.85)
    ax.scatter(steps[long_mask], padded_length[long_mask], color=VERMILLION, s=15, zorder=4, label="≥7,500 tokens")
    ax.axhline(8192, color=INK, linestyle=(0, (2, 2)), linewidth=0.8, alpha=0.7, label="8,192 limit")
    ax.scatter([507], [padded_length[step_507_index]], marker="*", color=ORANGE, edgecolor=INK, linewidth=0.35, s=58, zorder=6)
    ax.annotate(
        "long batch\n8,106 tokens",
        (507, padded_length[step_507_index]),
        xytext=(-54, -31),
        textcoords="offset points",
        arrowprops={"arrowstyle": "-", "color": GRAY, "linewidth": 0.7},
        fontsize=7,
        color=INK,
    )
    ax.set_title("Long-context exposure")
    ax.set_ylabel("Padded sequence length")
    ax.set_ylim(3000, 8400)
    ax.yaxis.set_major_locator(MaxNLocator(5))
    ax.legend(loc="lower right")
    panel_label(ax, "c")

    ax = axes[3]
    ax.plot(steps, throughput, color=SKY, alpha=0.30, linewidth=0.7, label="per-step")
    ax.plot(steps, rolling(throughput, 25), color=BLUE, linewidth=1.8, label="rolling median (25)")
    throughput_median = float(np.median(throughput))
    ax.axhline(throughput_median, color=INK, linestyle=(0, (2, 2)), linewidth=0.9)
    ax.text(700, throughput_median + 5, f"median {throughput_median:.1f}", fontsize=7, ha="right", color=INK)
    ax.set_title("Training throughput")
    ax.set_ylabel("Tokens / s / GPU")
    ax.legend(loc="lower left")
    panel_label(ax, "d")

    ax = axes[4]
    ax.plot(steps, step_time, color=LIGHT_GRAY, alpha=0.75, linewidth=0.7, label="step time")
    ax.plot(steps, rolling(step_time, 25), color=INK, linewidth=1.8, label="step rolling median")
    ax.scatter([507], [step_time[step_507_index]], color=VERMILLION, marker="*", s=50, zorder=5)
    ax.annotate("53.5 s", (507, step_time[step_507_index]), xytext=(7, -2), textcoords="offset points", fontsize=7, color=VERMILLION)
    ax2 = ax.twinx()
    ax2.plot(steps, rolling(optimizer_time, 25), color=GREEN, linewidth=1.25, linestyle=(0, (4, 2)), label="optimizer rolling median")
    ax2.set_ylabel("Optimizer step (s)", color=GREEN)
    ax2.tick_params(axis="y", colors=GREEN, labelsize=7.5)
    ax2.spines["right"].set_visible(True)
    ax2.spines["right"].set_color(GREEN)
    lines, labels = ax.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax.legend(lines + lines2, labels + labels2, loc="upper left")
    ax.set_title("Step-time decomposition")
    ax.set_ylabel("End-to-end step time (s)")
    panel_label(ax, "e")

    ax = axes[5]
    ax.plot(steps, padding_efficiency, color=ORANGE, alpha=0.25, linewidth=0.7, label="per-step efficiency")
    ax.plot(steps, rolling(padding_efficiency, 25), color=VERMILLION, linewidth=1.8, label="rolling median (25)")
    ax.set_ylim(0.45, 0.9)
    ax.set_ylabel("Non-padding token ratio")
    ax2 = ax.twinx()
    ax2.plot(steps, learning_rate * 1e5, color=PINK, linestyle=(0, (4, 2)), linewidth=1.25, label="cosine LR")
    ax2.set_ylabel("Learning rate (×10⁻⁵)", color=PINK)
    ax2.tick_params(axis="y", colors=PINK, labelsize=7.5)
    ax2.spines["right"].set_visible(True)
    ax2.spines["right"].set_color(PINK)
    lines, labels = ax.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax.legend(lines + lines2, labels + labels2, loc="lower right")
    ax.set_title(f"Packing efficiency and LR · {cumulative_tokens[-1] / 1e6:.2f}M tokens")
    panel_label(ax, "f")

    for ax in axes:
        ax.set_xlim(1, 710)
        ax.set_xlabel("Optimizer step")
        ax.xaxis.set_major_locator(MaxNLocator(6, integer=True))

    fig.suptitle(
        "CodePin-SFT · Qwen3.5-0.8B Full-Parameter Training Dynamics",
        fontsize=12.5,
        fontweight="bold",
        y=0.965,
    )
    fig.text(
        0.5,
        0.915,
        "4× V100 32GB · FSDP FP16 · global batch 8 · max context 8,192 · 710 steps",
        ha="center",
        va="top",
        fontsize=8.5,
        color=GRAY,
    )
    fig.savefig(output_dir / "fig_sft_training_dynamics.pdf")
    fig.savefig(output_dir / "fig_sft_training_dynamics.png", dpi=300)
    plt.close(fig)

    # Systems figure: expose the length-dependent cost across the continuous run.
    fig, axes = plt.subplots(1, 2, figsize=(7.4, 3.05), constrained_layout=True)
    for ax, y, ylabel, title in [
        (axes[0], step_time, "Step time (s)", "Sequence length drives wall time"),
        (axes[1], throughput, "Tokens / s / GPU", "Per-GPU throughput under variable packing"),
    ]:
        regular = ~long_mask
        ax.scatter(padded_length[regular], y[regular], s=11, color=SKY, alpha=0.45, edgecolors="none", label="all steps")
        ax.scatter(padded_length[long_mask], y[long_mask], s=18, color=ORANGE, alpha=0.78, edgecolors="none", label="≥7,500 tokens")
        slope, intercept = np.polyfit(padded_length, y, 1)
        xline = np.linspace(float(padded_length.min()), float(padded_length.max()), 100)
        ax.plot(xline, slope * xline + intercept, color=BLUE, linewidth=1.7, label="linear fit")
        ax.scatter(
            [padded_length[step_507_index]],
            [y[step_507_index]],
            marker="*",
            s=65,
            color=VERMILLION,
            edgecolor=INK,
            linewidth=0.35,
            zorder=6,
            label="step 507",
        )
        ax.axvspan(7500, 8192, color=VERMILLION, alpha=0.06, linewidth=0)
        ax.set_xlabel("Padded sequence length")
        ax.set_ylabel(ylabel)
        ax.set_title(title)
        ax.xaxis.set_major_formatter(FuncFormatter(lambda value, _: f"{value / 1000:.0f}k"))
        ax.legend(loc="best")
    panel_label(axes[0], "a")
    panel_label(axes[1], "b")
    fig.suptitle("Length-Dependent Training Behavior on 4× V100", fontsize=11.5, fontweight="bold")
    fig.savefig(output_dir / "fig_sft_length_efficiency.pdf")
    fig.savefig(output_dir / "fig_sft_length_efficiency.png", dpi=300)
    plt.close(fig)

    print(json.dumps({
        "outputs": [
            "fig_sft_training_dynamics.pdf",
            "fig_sft_training_dynamics.png",
            "fig_sft_length_efficiency.pdf",
            "fig_sft_length_efficiency.png",
            "sft_training_metrics.csv",
            "sft_training_summary.json",
        ],
        "summary": summary,
    }, indent=2))


if __name__ == "__main__":
    main()