"""Generate publication figures from results CSVs."""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def fig_optimality_gap_bar(df: pd.DataFrame, out: Path) -> None:
    agg = df.groupby("method")["optimality_gap_pct"].mean().sort_values()
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.barh(agg.index, agg.values)
    ax.set_xlabel("Mean optimality gap (%)")
    ax.set_title("Mean optimality gap per method (lower is better)")
    ax.grid(True, alpha=0.3, axis="x")
    plt.tight_layout()
    plt.savefig(out, dpi=150)
    plt.close()


def fig_per_family_gap(df: pd.DataFrame, out: Path) -> None:
    pivot = df.groupby(["family", "method"])["optimality_gap_pct"].mean().unstack()
    fig, ax = plt.subplots(figsize=(10, 6))
    pivot.plot(kind="bar", ax=ax)
    ax.set_ylabel("Mean optimality gap (%)")
    ax.set_title("Per-family optimality gap")
    ax.grid(True, alpha=0.3, axis="y")
    plt.tight_layout()
    plt.savefig(out, dpi=150)
    plt.close()


def fig_pareto(df: pd.DataFrame, out: Path) -> None:
    if "total_tardiness" not in df.columns:
        return
    means = df.groupby("method").agg(
        mk=("makespan", "mean"),
        td=("total_tardiness", "mean"),
    ).reset_index()
    fig, ax = plt.subplots(figsize=(8, 6))
    for _, r in means.iterrows():
        ax.scatter(r["mk"], r["td"], s=80)
        ax.annotate(r["method"], (r["mk"], r["td"]),
                    xytext=(5, 5), textcoords="offset points", fontsize=8)
    ax.set_xlabel("Mean makespan")
    ax.set_ylabel("Mean tardiness")
    ax.set_title("Pareto: makespan vs tardiness")
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(out, dpi=150)
    plt.close()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--results", default="results/benchmark_full.csv")
    parser.add_argument("--out-dir", default="results/figures")
    args = parser.parse_args()

    df = pd.read_csv(args.results)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    fig_optimality_gap_bar(df, out_dir / "optimality_gap_bar.png")
    fig_per_family_gap(df, out_dir / "per_family_gap.png")
    fig_pareto(df, out_dir / "pareto.png")
    print(f"Wrote figures to {out_dir}/")


if __name__ == "__main__":
    main()