#!/usr/bin/env python3
"""Plot training loss curves from training logs for comparison."""

import re
import sys
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

def parse_loss_from_log(log_path):
    """Extract (step, loss, grad_norm) from training log."""
    steps, losses, grad_norms = [], [], []
    pattern = re.compile(r"'loss': ([\d.]+), 'grad_norm': ([\d.]+)")
    with open(log_path) as f:
        for line in f:
            m = pattern.search(line)
            if m:
                steps.append(len(losses) + 1)
                losses.append(float(m.group(1)))
                grad_norms.append(float(m.group(2)))
    return steps, losses, grad_norms


def parse_loss_from_tensorboard(tb_dir):
    """Extract (step, loss, grad_norm) from tensorboard event files."""
    try:
        from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
    except ImportError:
        print("tensorboard not installed, skipping TB parsing")
        return [], [], []

    ea = EventAccumulator(tb_dir)
    ea.Reload()

    steps, losses, grad_norms = [], [], []
    if 'train/loss' in ea.scalars.Keys():
        for event in ea.scalars.Items('train/loss'):
            steps.append(event.step)
            losses.append(event.value)
    elif 'loss' in ea.scalars.Keys():
        for event in ea.scalars.Items('loss'):
            steps.append(event.step)
            losses.append(event.value)

    if 'train/grad_norm' in ea.scalars.Keys():
        for event in ea.scalars.Items('train/grad_norm'):
            grad_norms.append(event.value)
    elif 'grad_norm' in ea.scalars.Keys():
        for event in ea.scalars.Items('grad_norm'):
            grad_norms.append(event.value)

    return steps, losses, grad_norms


def plot_comparison(curves, output_path, title="Training Loss Comparison"):
    """Plot loss and grad_norm for multiple runs."""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    for name, (steps, losses, grad_norms) in curves.items():
        if losses:
            ax1.plot(steps[:len(losses)], losses, label=name, alpha=0.8)
        if grad_norms:
            ax2.plot(steps[:len(grad_norms)], grad_norms, label=name, alpha=0.8)

    ax1.set_xlabel('Step')
    ax1.set_ylabel('Loss')
    ax1.set_title('Training Loss')
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    ax2.set_xlabel('Step')
    ax2.set_ylabel('Grad Norm')
    ax2.set_title('Gradient Norm')
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    plt.suptitle(title, fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"Saved: {output_path}")


if __name__ == '__main__':
    import os
    import glob

    curves = {}

    # Baseline: ddp-verify (from tensorboard)
    baseline_tb_dirs = sorted(glob.glob('/mnt/bn/leonworkspace/terry/model/qwen3vl-4b-roi-K24T3-185k-ddp-verify/runs/*'))
    if baseline_tb_dirs:
        steps, losses, grad_norms = parse_loss_from_tensorboard(baseline_tb_dirs[-1])
        if losses:
            curves['baseline (ddp-verify)'] = (steps, losses, grad_norms)

    # Bidir: from live log
    bidir_log = '/tmp/bidir_train.log'
    if os.path.exists(bidir_log):
        steps, losses, grad_norms = parse_loss_from_log(bidir_log)
        if losses:
            curves['bidir (all_visual)'] = (steps, losses, grad_norms)

    # Bidir: from tensorboard (if available)
    bidir_tb_dirs = sorted(glob.glob('/mnt/bn/leonworkspace/terry/model/qwen3vl-4b-roi-K24T3-185k-bidir/runs/*'))
    if bidir_tb_dirs:
        steps, losses, grad_norms = parse_loss_from_tensorboard(bidir_tb_dirs[-1])
        if losses and len(losses) > 10:
            curves['bidir (tensorboard)'] = (steps, losses, grad_norms)

    if not curves:
        print("No data found!")
        sys.exit(1)

    output_path = '/opt/tiger/thothvl_pretrain/visualize/bidir_vs_baseline_loss.png'
    plot_comparison(curves, output_path, "Bidir Visual Attention Training: Loss Comparison")