from __future__ import annotations

"""Generate training/validation metric plots from CSV logs.

Reads train.csv and val.csv, produces publication-quality plots
using seaborn + matplotlib. Called automatically after each eval epoch
or standalone via CLI.

Usage:
    python -m src.training.plot_metrics --log-dir out/gtauav/with_text/logs
"""

import argparse
import logging
from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns

LOGGER = logging.getLogger("caption_test.plot")

# Seaborn style.
sns.set_theme(style="whitegrid", palette="deep", font_scale=1.1)


def plot_train_metrics(train_df: pd.DataFrame, out_dir: Path) -> None:
    """Plot training metrics: loss, recall, AP, temperature, gates, lr."""
    fig, axes = plt.subplots(2, 3, figsize=(20, 10))
    fig.suptitle("Training Metrics", fontsize=16, fontweight="bold")

    # 1. Loss.
    ax = axes[0, 0]
    col = "train_loss" if "train_loss" in train_df.columns else "total"
    if col in train_df.columns:
        sns.lineplot(data=train_df, x="epoch", y=col, ax=ax, marker="o", linewidth=2)
    ax.set_title("Train Loss")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("InfoNCE Loss")

    # 2. Recall@K (train).
    ax = axes[0, 1]
    for k in [1, 5, 10]:
        col = f"r@{k}_q2g"
        if col in train_df.columns:
            df_valid = train_df.dropna(subset=[col])
            if not df_valid.empty:
                sns.lineplot(data=df_valid, x="epoch", y=col, ax=ax, marker="o", linewidth=2, label=f"R@{k}")
    ax.set_title("Train Recall@K (drone → sat)")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Recall")
    ax.set_ylim(0, 1)
    ax.legend()

    # 3. Average Precision (train).
    ax = axes[0, 2]
    if "ap_q2g" in train_df.columns:
        df_valid = train_df.dropna(subset=["ap_q2g"])
        if not df_valid.empty:
            sns.lineplot(data=df_valid, x="epoch", y="ap_q2g", ax=ax, marker="s", linewidth=2, color="purple")
    ax.set_title("Train AP (drone → sat)")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("AP")
    ax.set_ylim(0, 1)

    # 4. Temperature (tau).
    ax = axes[1, 0]
    if "temperature" in train_df.columns:
        sns.lineplot(data=train_df, x="epoch", y="temperature", ax=ax, marker="s", linewidth=2, color="orange")
    ax.set_title("Temperature (τ)")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("τ")

    # 5. Gate values.
    ax = axes[1, 1]
    if "gate_q" in train_df.columns and "gate_g" in train_df.columns:
        sns.lineplot(data=train_df, x="epoch", y="gate_q", ax=ax, marker="o", linewidth=2, label="gate_q (drone)")
        sns.lineplot(data=train_df, x="epoch", y="gate_g", ax=ax, marker="s", linewidth=2, label="gate_g (sat)")
        ax.legend()
    ax.set_title("Gate Values (σ(α))")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Image weight")
    ax.set_ylim(0, 1)

    # 6. Learning rate.
    ax = axes[1, 2]
    if "lr" in train_df.columns:
        sns.lineplot(data=train_df, x="epoch", y="lr", ax=ax, marker="^", linewidth=2, color="green")
    ax.set_title("Learning Rate")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("LR")
    ax.ticklabel_format(axis="y", style="scientific", scilimits=(0, 0))

    plt.tight_layout()
    path = out_dir / "train_metrics.png"
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    LOGGER.info("📊 Saved %s", path)


def plot_val_metrics(val_df: pd.DataFrame, out_dir: Path, train_recall_df: pd.DataFrame | None = None) -> None:
    """Plot recall + AP metrics: train vs val."""
    fig, axes = plt.subplots(1, 3, figsize=(20, 5))
    fig.suptitle("Retrieval Metrics (train vs val)", fontsize=16, fontweight="bold")

    # 1. Recall@K (q→g): train + val.
    ax = axes[0]
    for k in [1, 5, 10]:
        col = f"r@{k}_q2g"
        if col in val_df.columns:
            sns.lineplot(data=val_df, x="epoch", y=col, ax=ax, marker="o", linewidth=2, label=f"val R@{k}")
        if train_recall_df is not None and col in train_recall_df.columns:
            sns.lineplot(data=train_recall_df, x="epoch", y=col, ax=ax, marker="x", linewidth=1.5, linestyle="--", label=f"train R@{k}")
    ax.set_title("Recall@K (drone → satellite)")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Recall")
    ax.set_ylim(0, 1)
    ax.legend(fontsize=8)

    # 2. Recall@K (g→q): train + val.
    ax = axes[1]
    for k in [1, 5, 10]:
        col = f"r@{k}_g2q"
        if col in val_df.columns:
            sns.lineplot(data=val_df, x="epoch", y=col, ax=ax, marker="s", linewidth=2, label=f"val R@{k}")
        if train_recall_df is not None and col in train_recall_df.columns:
            sns.lineplot(data=train_recall_df, x="epoch", y=col, ax=ax, marker="x", linewidth=1.5, linestyle="--", label=f"train R@{k}")
    ax.set_title("Recall@K (satellite → drone)")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Recall")
    ax.set_ylim(0, 1)
    ax.legend(fontsize=8)

    # 3. AP (train vs val, both directions).
    ax = axes[2]
    if "ap_q2g" in val_df.columns:
        sns.lineplot(data=val_df, x="epoch", y="ap_q2g", ax=ax, marker="o", linewidth=2, color="royalblue", label="val AP q→g")
    if "ap_g2q" in val_df.columns:
        sns.lineplot(data=val_df, x="epoch", y="ap_g2q", ax=ax, marker="s", linewidth=2, color="coral", label="val AP g→q")
    if train_recall_df is not None:
        if "ap_q2g" in train_recall_df.columns:
            sns.lineplot(data=train_recall_df, x="epoch", y="ap_q2g", ax=ax, marker="x", linewidth=1.5, linestyle="--", color="royalblue", label="train AP q→g")
        if "ap_g2q" in train_recall_df.columns:
            sns.lineplot(data=train_recall_df, x="epoch", y="ap_g2q", ax=ax, marker="x", linewidth=1.5, linestyle="--", color="coral", label="train AP g→q")
    ax.set_title("Average Precision")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("AP")
    ax.set_ylim(0, 1)
    ax.legend(fontsize=8)

    plt.tight_layout()
    path = out_dir / "val_metrics.png"
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    LOGGER.info("📊 Saved %s", path)


def plot_combined(train_df: pd.DataFrame, val_df: pd.DataFrame, out_dir: Path) -> None:
    """Combined overview: loss + R@1 on same figure."""
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    fig.suptitle("Training Overview", fontsize=16, fontweight="bold")

    # 1. Train + Val loss.
    ax = axes[0]
    sns.lineplot(data=train_df, x="epoch", y="total", ax=ax, marker="o", linewidth=2, color="crimson", label="train")
    if "loss" in val_df.columns:
        sns.lineplot(data=val_df, x="epoch", y="loss", ax=ax, marker="s", linewidth=2, color="royalblue", label="val")
    ax.set_title("Loss (train vs val)")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("InfoNCE Loss")
    ax.legend()

    # 2. Val R@1 q→g.
    ax = axes[1]
    if "r@1_q2g" in val_df.columns:
        sns.lineplot(data=val_df, x="epoch", y="r@1_q2g", ax=ax, marker="o", linewidth=2, color="royalblue", label="R@1 q→g")
    if "r@1_g2q" in val_df.columns:
        sns.lineplot(data=val_df, x="epoch", y="r@1_g2q", ax=ax, marker="s", linewidth=2, color="coral", label="R@1 g→q")
    ax.set_title("Validation R@1")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Recall@1")
    ax.set_ylim(0, None)
    ax.legend()

    # 3. Gates + Temperature.
    ax = axes[2]
    if "gate_q" in train_df.columns:
        sns.lineplot(data=train_df, x="epoch", y="gate_q", ax=ax, linewidth=2, label="gate_q")
        sns.lineplot(data=train_df, x="epoch", y="gate_g", ax=ax, linewidth=2, label="gate_g")
    ax2 = ax.twinx()
    sns.lineplot(data=train_df, x="epoch", y="temperature", ax=ax2, linewidth=2, color="orange", linestyle="--", label="τ")
    ax.set_title("Gates & Temperature")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Gate value")
    ax2.set_ylabel("τ")
    lines1, labels1 = ax.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax.legend(lines1 + lines2, labels1 + labels2, loc="best")
    ax2.get_legend().remove() if ax2.get_legend() else None

    plt.tight_layout()
    path = out_dir / "overview.png"
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    LOGGER.info("📊 Saved %s", path)


def generate_plots(log_dir: str | Path) -> None:
    """Generate all plots from CSV logs in log_dir."""
    log_dir = Path(log_dir)

    train_csv = log_dir / "train.csv"
    val_csv = log_dir / "val.csv"

    if train_csv.exists():
        train_df = pd.read_csv(train_csv)
        plot_train_metrics(train_df, log_dir)
    else:
        LOGGER.warning("No train.csv found in %s", log_dir)
        train_df = None

    train_recall_csv = log_dir / "train_recall.csv"
    train_recall_df = None
    if train_recall_csv.exists():
        train_recall_df = pd.read_csv(train_recall_csv)

    if val_csv.exists():
        val_df = pd.read_csv(val_csv)
        plot_val_metrics(val_df, log_dir, train_recall_df=train_recall_df)
    else:
        LOGGER.warning("No val.csv found in %s", log_dir)
        val_df = None

    if train_df is not None and val_df is not None:
        plot_combined(train_df, val_df, log_dir)


def main() -> None:
    import coloredlogs
    coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")

    parser = argparse.ArgumentParser(description="Generate metric plots from CSV logs.")
    parser.add_argument("--log-dir", type=str, required=True, help="Path to logs/ directory.")
    args = parser.parse_args()
    generate_plots(args.log_dir)


if __name__ == "__main__":
    main()