#!/usr/bin/env python3
# plot_training.py -- plot training metrics from trainer_state.json
# Usage: python plot_training.py [path/to/checkpoint-XXXX]
#        python plot_training.py  (auto-finds latest checkpoint)

import json
import sys
import glob
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

BG = "#0D1117"
CARD = "#161B22"
BORDER = "#30363D"
TEXT = "#E6EDF3"
SUBTLE = "#8B949E"
GRID = "#21262D"
C_TRAIN = "#79C0FF"
C_EVAL = "#56D364"
C_LR = "#FFA657"

def find_latest_checkpoint():
    checkpoints = sorted(glob.glob("checkpoints/checkpoint-*"))
    if not checkpoints:
        print("ERROR: No checkpoints found in ./checkpoints/")
        sys.exit(1)
    return checkpoints[-1]

def load_logs(checkpoint_dir):
    state_file = os.path.join(checkpoint_dir, "trainer_state.json")
    if not os.path.exists(state_file):
        print(f"ERROR: {state_file} not found")
        sys.exit(1)
    with open(state_file) as f:
        state = json.load(f)
    return state["log_history"]

def plot(log_history, outpath="training_plot.png"):
    steps_train, loss_train = [], []
    steps_eval, loss_eval = [], []
    steps_lr, lr_vals = [], []

    for entry in log_history:
        if "loss" in entry:
            steps_train.append(entry["step"])
            loss_train.append(entry["loss"])
        if "eval_loss" in entry:
            steps_eval.append(entry["step"])
            loss_eval.append(entry["eval_loss"])
        if "learning_rate" in entry:
            steps_lr.append(entry["step"])
            lr_vals.append(entry["learning_rate"])

    fig, ax1 = plt.subplots(figsize=(14, 8))
    fig.patch.set_facecolor(BG)
    ax1.set_facecolor(CARD)

    ax1.plot(steps_train, loss_train, color=C_TRAIN, alpha=0.4,
             linewidth=1, label="Train loss")
    ax1.plot(steps_eval, loss_eval, color=C_EVAL, linewidth=2.5,
             marker="o", markersize=7, label="Eval loss")

    ax1.set_xlabel("Step", fontsize=13, color=TEXT)
    ax1.set_ylabel("Loss", fontsize=13, color=TEXT)
    ax1.tick_params(colors=SUBTLE)
    ax1.grid(True, alpha=0.15, color=GRID, linestyle="--")
    for spine in ax1.spines.values():
        spine.set_color(BORDER)

    ax2 = ax1.twinx()
    ax2.set_facecolor(CARD)
    ax2.plot(steps_lr, lr_vals, color=C_LR, linewidth=1.5, alpha=0.7,
             linestyle="--", label="Learning rate")
    ax2.set_ylabel("Learning rate", fontsize=13, color=C_LR)
    ax2.tick_params(axis="y", colors=C_LR)
    for spine in ax2.spines.values():
        spine.set_color(BORDER)

    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2,
               loc="upper right", fontsize=11, framealpha=0.92,
               facecolor=CARD, edgecolor=BORDER, labelcolor=TEXT)

    n_steps = max(steps_train) if steps_train else 0
    ax1.set_title(f"Training Metrics ({n_steps} steps)",
                 fontsize=16, fontweight="bold", color=TEXT, pad=20)

    fig.text(0.5, 0.02,
             f"LFM2.5-VL-450M hand-tracking | LoRA r=16 | {len(loss_train)} train + {len(loss_eval)} eval points",
             ha="center", fontsize=10, color=SUBTLE, style="italic")

    plt.tight_layout(rect=[0, 0.06, 1, 1])
    fig.savefig(outpath, dpi=150, bbox_inches="tight",
                facecolor=fig.get_facecolor(), edgecolor="none")
    plt.close(fig)
    print(f"Saved to {outpath}")

if __name__ == "__main__":
    ckpt = sys.argv[1] if len(sys.argv) > 1 else find_latest_checkpoint()
    print(f"Loading from {ckpt}")
    logs = load_logs(ckpt)
    print(f"  {len(logs)} log entries")
    out = sys.argv[2] if len(sys.argv) > 2 else "training_plot.png"
    plot(logs, out)