"""
Gera graficos de treinamento do ResearcherBot-1 Multimodal.

Produz 4 graficos:
1. Training & Validation Loss
2. Perplexity (Train & Val)
3. Learning Rate Schedule
4. MoE Auxiliary Loss (Validation)

Salva em: training_plots/
"""

import json
import os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker

# Estilo escuro profissional
plt.rcParams.update({
    "figure.facecolor": "#0d1117",
    "axes.facecolor": "#161b22",
    "axes.edgecolor": "#30363d",
    "axes.labelcolor": "#c9d1d9",
    "xtick.color": "#8b949e",
    "ytick.color": "#8b949e",
    "text.color": "#c9d1d9",
    "grid.color": "#21262d",
    "grid.linestyle": "--",
    "grid.alpha": 0.6,
    "font.family": "sans-serif",
    "font.size": 11,
    "legend.facecolor": "#161b22",
    "legend.edgecolor": "#30363d",
    "legend.fontsize": 10,
})

COLORS = {
    "train": "#58a6ff",
    "val": "#f78166",
    "lr": "#7ee787",
    "aux": "#d2a8ff",
    "accent": "#ffa657",
}

OUTPUT_DIR = "training_plots"


def load_metrics(path="training_metrics.json"):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def plot_loss(m, save_dir):
    """Grafico 1: Train & Val Loss."""
    fig, ax = plt.subplots(figsize=(10, 5.5))

    ax.plot(m["train_steps"], m["train_loss"], color=COLORS["train"],
            linewidth=1.8, label="Train Loss", alpha=0.9)
    ax.plot(m["val_steps"], m["val_loss"], color=COLORS["val"],
            linewidth=2.2, marker="o", markersize=7, label="Val Loss",
            linestyle="--", alpha=0.9)

    ax.set_xlabel("Step")
    ax.set_ylabel("Loss")
    ax.set_title("Training & Validation Loss", fontsize=14, fontweight="bold", pad=12)
    ax.legend(loc="upper right")
    ax.grid(True)
    ax.set_ylim(bottom=0)

    # Anotacao do melhor val loss
    best_idx = m["val_loss"].index(min(m["val_loss"]))
    best_step = m["val_steps"][best_idx]
    best_val = m["val_loss"][best_idx]
    ax.annotate(f"Best: {best_val:.4f}",
                xy=(best_step, best_val),
                xytext=(best_step + 100, best_val + 0.5),
                arrowprops=dict(arrowstyle="->", color=COLORS["accent"], lw=1.5),
                fontsize=10, color=COLORS["accent"], fontweight="bold")

    plt.tight_layout()
    path = os.path.join(save_dir, "loss.png")
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Salvo: {path}")


def plot_perplexity(m, save_dir):
    """Grafico 2: Train & Val Perplexity."""
    fig, ax = plt.subplots(figsize=(10, 5.5))

    ax.plot(m["train_steps"], m["train_ppl"], color=COLORS["train"],
            linewidth=1.8, label="Train PPL", alpha=0.9)
    ax.plot(m["val_steps"], m["val_ppl"], color=COLORS["val"],
            linewidth=2.2, marker="s", markersize=7, label="Val PPL",
            linestyle="--", alpha=0.9)

    ax.set_xlabel("Step")
    ax.set_ylabel("Perplexity")
    ax.set_title("Perplexity (Train & Validation)", fontsize=14, fontweight="bold", pad=12)
    ax.legend(loc="upper right")
    ax.grid(True)
    ax.set_yscale("log")

    # Anotacao PPL final
    final_train_ppl = m["train_ppl"][-1]
    ax.axhline(y=final_train_ppl, color=COLORS["train"], linestyle=":", alpha=0.4)
    ax.text(m["train_steps"][-1], final_train_ppl * 1.1,
            f"Final: {final_train_ppl:.1f}",
            color=COLORS["train"], fontsize=9, ha="right")

    plt.tight_layout()
    path = os.path.join(save_dir, "perplexity.png")
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Salvo: {path}")


def plot_learning_rate(m, save_dir):
    """Grafico 3: Learning Rate Schedule."""
    fig, ax = plt.subplots(figsize=(10, 5.5))

    ax.plot(m["train_steps"], m["lr"], color=COLORS["lr"],
            linewidth=2.0, alpha=0.9)
    ax.fill_between(m["train_steps"], 0, m["lr"],
                     color=COLORS["lr"], alpha=0.15)

    ax.set_xlabel("Step")
    ax.set_ylabel("Learning Rate")
    ax.set_title("Learning Rate Schedule (Warmup + Cosine Decay)",
                 fontsize=14, fontweight="bold", pad=12)
    ax.grid(True)
    ax.yaxis.set_major_formatter(mticker.FormatStrFormatter("%.1e"))

    # Marca warmup
    warmup_end = 500
    ax.axvline(x=warmup_end, color=COLORS["accent"], linestyle="--",
               alpha=0.6, linewidth=1.5)
    ax.text(warmup_end + 20, max(m["lr"]) * 0.95, "Warmup\nEnd",
            color=COLORS["accent"], fontsize=9, fontweight="bold")

    plt.tight_layout()
    path = os.path.join(save_dir, "learning_rate.png")
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Salvo: {path}")


def plot_aux_loss(m, save_dir):
    """Grafico 4: MoE Auxiliary Loss."""
    if not m.get("val_aux_loss"):
        return

    fig, ax = plt.subplots(figsize=(10, 5.5))

    ax.plot(m["val_steps"], m["val_aux_loss"], color=COLORS["aux"],
            linewidth=2.2, marker="D", markersize=7, alpha=0.9)
    ax.fill_between(m["val_steps"], 0, m["val_aux_loss"],
                     color=COLORS["aux"], alpha=0.1)

    ax.set_xlabel("Step")
    ax.set_ylabel("Aux Loss (Load Balance)")
    ax.set_title("MoE Load Balancing Loss (Validation)",
                 fontsize=14, fontweight="bold", pad=12)
    ax.grid(True)

    # Ideal = 1.0 (uniforme)
    ax.axhline(y=1.0, color="#8b949e", linestyle=":", alpha=0.5, linewidth=1.2)
    ax.text(m["val_steps"][-1], 1.05, "Ideal (uniforme)",
            color="#8b949e", fontsize=9, ha="right")

    plt.tight_layout()
    path = os.path.join(save_dir, "moe_aux_loss.png")
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Salvo: {path}")


def plot_combined_dashboard(m, save_dir):
    """Grafico 5: Dashboard combinado 2x2."""
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))

    # (0,0) Loss
    ax = axes[0, 0]
    ax.plot(m["train_steps"], m["train_loss"], color=COLORS["train"],
            linewidth=1.5, label="Train", alpha=0.9)
    ax.plot(m["val_steps"], m["val_loss"], color=COLORS["val"],
            linewidth=2, marker="o", markersize=5, label="Val",
            linestyle="--", alpha=0.9)
    ax.set_xlabel("Step")
    ax.set_ylabel("Loss")
    ax.set_title("Loss", fontsize=12, fontweight="bold")
    ax.legend(fontsize=9)
    ax.grid(True)
    ax.set_ylim(bottom=0)

    # (0,1) PPL
    ax = axes[0, 1]
    ax.plot(m["train_steps"], m["train_ppl"], color=COLORS["train"],
            linewidth=1.5, label="Train", alpha=0.9)
    ax.plot(m["val_steps"], m["val_ppl"], color=COLORS["val"],
            linewidth=2, marker="s", markersize=5, label="Val",
            linestyle="--", alpha=0.9)
    ax.set_xlabel("Step")
    ax.set_ylabel("Perplexity")
    ax.set_title("Perplexity", fontsize=12, fontweight="bold")
    ax.legend(fontsize=9)
    ax.grid(True)
    ax.set_yscale("log")

    # (1,0) LR
    ax = axes[1, 0]
    ax.plot(m["train_steps"], m["lr"], color=COLORS["lr"], linewidth=1.8, alpha=0.9)
    ax.fill_between(m["train_steps"], 0, m["lr"], color=COLORS["lr"], alpha=0.1)
    ax.set_xlabel("Step")
    ax.set_ylabel("LR")
    ax.set_title("Learning Rate", fontsize=12, fontweight="bold")
    ax.grid(True)
    ax.yaxis.set_major_formatter(mticker.FormatStrFormatter("%.1e"))

    # (1,1) Aux Loss
    ax = axes[1, 1]
    if m.get("val_aux_loss"):
        ax.plot(m["val_steps"], m["val_aux_loss"], color=COLORS["aux"],
                linewidth=2, marker="D", markersize=5, alpha=0.9)
        ax.axhline(y=1.0, color="#8b949e", linestyle=":", alpha=0.4)
    ax.set_xlabel("Step")
    ax.set_ylabel("Aux Loss")
    ax.set_title("MoE Aux Loss", fontsize=12, fontweight="bold")
    ax.grid(True)

    fig.suptitle("ResearcherBot-1 Multimodal  |  19.9M params  |  MoE 8 experts top-2",
                 fontsize=15, fontweight="bold", y=0.98, color="#58a6ff")
    plt.tight_layout(rect=[0, 0, 1, 0.96])
    path = os.path.join(save_dir, "dashboard.png")
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Salvo: {path}")


def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print("=" * 50)
    print("  Gerando graficos de treinamento...")
    print("=" * 50)

    m = load_metrics()
    print(f"  Steps: {m['total_steps']} | Tempo: {m['total_time_min']:.1f} min")
    print(f"  Best Val Loss: {m['best_val_loss']:.4f}")
    print()

    plot_loss(m, OUTPUT_DIR)
    plot_perplexity(m, OUTPUT_DIR)
    plot_learning_rate(m, OUTPUT_DIR)
    plot_aux_loss(m, OUTPUT_DIR)
    plot_combined_dashboard(m, OUTPUT_DIR)

    print(f"\n  Todos os graficos em: {OUTPUT_DIR}/")


if __name__ == "__main__":
    main()