"""
Generate all paper figures from measurement CSVs.
"""

import argparse
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


def plot_lpips_bar(df, output_dir):
    order = df.groupby("config")["lpips_mean"].first().sort_values().index
    means = df.set_index("config").reindex(order)

    fig, ax = plt.subplots(figsize=(12, 5))
    colors = plt.cm.RdYlGn_r(np.linspace(0.2, 0.9, len(order)))
    ax.bar(range(len(order)), means["lpips_mean"], yerr=means["lpips_std"],
           capsize=4, color=colors)
    ax.set_xticks(range(len(order)))
    ax.set_xticklabels(order, rotation=45, ha="right", fontsize=9)
    ax.set_ylabel("LPIPS (vs Q8_0)")
    ax.set_title("Perceptual Distance from Q8_0 Baseline")
    ax.grid(axis="y", alpha=0.3)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "lpips_bar.png"), dpi=200)
    plt.close()


def plot_lpips_heatmap(df, output_dir):
    cats = ["A", "B", "C", "D", "E"]
    cat_cols = [c for c in df.columns if c.startswith("lpips_") and len(c) == 7]
    if not cat_cols:
        return

    order = df.sort_values("lpips_mean")["config"]
    pivot = df.set_index("config").reindex(order)[cat_cols]
    pivot.columns = [c.replace("lpips_", "Cat ") for c in cat_cols]

    fig, ax = plt.subplots(figsize=(10, max(4, len(order)*0.5)))
    im = ax.imshow(pivot.values, cmap="YlOrRd", aspect="auto")
    ax.set_xticks(range(len(pivot.columns)))
    ax.set_xticklabels(pivot.columns)
    ax.set_yticks(range(len(order)))
    ax.set_yticklabels(order, fontsize=9)
    for i in range(pivot.shape[0]):
        for j in range(pivot.shape[1]):
            v = pivot.values[i, j]
            if not np.isnan(v):
                ax.text(j, i, f"{v:.3f}", ha="center", va="center", fontsize=8)
    plt.colorbar(im, ax=ax, label="LPIPS")
    ax.set_title("LPIPS by Category × Config")
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "lpips_heatmap.png"), dpi=200)
    plt.close()


def plot_quality_vs_lpips(lpips_df, qual_df, output_dir):
    merged = lpips_df.merge(qual_df, on="config", suffixes=("_lp", "_q"))

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    for ax, metric, label in [
        (axes[0], "quality_mean", "Quality Score"),
        (axes[1], "clip_mean_q" if "clip_mean_q" in merged.columns else "clip_mean", "CLIP Score"),
    ]:
        if metric not in merged.columns:
            continue
        uniform = merged[merged["config"].str.startswith("uniform")]
        asym = merged[merged["config"].str.startswith("asym")]

        ax.scatter(uniform["lpips_mean"], uniform[metric], c="steelblue",
                   s=80, label="Uniform", zorder=3)
        ax.scatter(asym["lpips_mean"], asym[metric], c="darkorange",
                   marker="D", s=80, label="Asymmetric", zorder=3)

        for _, row in merged.iterrows():
            ax.annotate(row["config"].replace("uniform_", "").replace("asym_", ""),
                        (row["lpips_mean"], row[metric]),
                        fontsize=7, alpha=0.7, textcoords="offset points",
                        xytext=(5, 5))
        ax.set_xlabel("LPIPS vs Q8_0")
        ax.set_ylabel(label)
        ax.legend()
        ax.grid(alpha=0.3)

    plt.suptitle("Reference-Based vs Reference-Free Metrics", fontsize=13)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "scatter_quality_vs_lpips.png"), dpi=200)
    plt.close()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--metrics-dir", required=True)
    parser.add_argument("--output-dir", required=True)
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    # Load summaries
    lp_path = os.path.join(args.metrics_dir, "summary_lpips_clip.csv")
    q_path = os.path.join(args.metrics_dir, "summary_quality_reffree.csv")

    if os.path.exists(lp_path):
        lp_df = pd.read_csv(lp_path)
        plot_lpips_bar(lp_df, args.output_dir)
        plot_lpips_heatmap(lp_df, args.output_dir)
        print("Generated: lpips_bar.png, lpips_heatmap.png")

        if os.path.exists(q_path):
            q_df = pd.read_csv(q_path)
            plot_quality_vs_lpips(lp_df, q_df, args.output_dir)
            print("Generated: scatter_quality_vs_lpips.png")

    print("Done!")


if __name__ == "__main__":
    main()