#!/usr/bin/env python3
from __future__ import annotations

import json
from pathlib import Path

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


WNS_KEYS = [
    "timing__setup__wns",
    "timing__setup__WNS",
    "finish__timing__setup__wns",
    "timing__setup__ws",
    "finish__timing__setup__ws",
    "route__timing__setup__ws",
    "cts__timing__setup__ws",
    "detailedplace__timing__setup__ws",
    "floorplan__timing__setup__ws",
    "globalplace__timing__setup__ws",
    "globalroute__timing__setup__ws",
    "placeopt__timing__setup__ws",
]

AREA_KEYS = [
    "finish__design__die__area",
    "globalroute__design__die__area",
    "placeopt__design__die__area",
    "detailedplace__design__die__area",
    "floorplan__design__die__area",
    "design__die__area",
]


def pick_first(metrics: dict, keys: list[str]) -> float | None:
    for k in keys:
        if k in metrics:
            try:
                return float(metrics[k])
            except Exception:
                return None
    lower = {kk.lower(): kk for kk in metrics.keys()}
    for k in keys:
        kk = lower.get(k.lower())
        if kk:
            try:
                return float(metrics[kk])
            except Exception:
                return None
    return None


def load_metrics(mj: str | None) -> dict | None:
    if not mj:
        return None
    try:
        return json.loads(mj)
    except Exception:
        return None


def main() -> None:
    import argparse

    p = argparse.ArgumentParser(description="Generate Pareto scatter (area vs WNS)")
    p.add_argument("--db", default="runs/sweep_finish.sqlite", help="SQLite database path")
    p.add_argument("--out", default="runs/sweep_finish/plots/pareto_area_vs_wns", help="Output path without extension")
    args = p.parse_args()

    con = sqlite3.connect(args.db)
    df = pd.read_sql_query("SELECT * FROM trials", con)
    con.close()

    rows = []
    for _, row in df.iterrows():
        metrics = load_metrics(row.get("metrics_json"))
        if not metrics:
            continue
        wns = pick_first(metrics, WNS_KEYS)
        area = pick_first(metrics, AREA_KEYS)
        if wns is None or area is None:
            continue
        rows.append({"wns": wns, "area": area, "fidelity": row.get("fidelity", "unknown")})

    if not rows:
        print("No points with both WNS and area found.")
        return

    plot_df = pd.DataFrame(rows)
    colors = {
        "synth": "#999999",
        "place": "#4B8BBE",
        "route": "#306998",
        "finish": "#E07B39",
        "unknown": "#666666",
    }
    markers = {
        "synth": "o",
        "place": "s",
        "route": "^",
        "finish": "D",
        "unknown": "o",
    }

    fig, ax = plt.subplots(figsize=(6.4, 4.6))
    for fid, group in plot_df.groupby("fidelity"):
        ax.scatter(
            group["area"],
            group["wns"],
            label=fid,
            color=colors.get(fid, "#666666"),
            marker=markers.get(fid, "o"),
            alpha=0.75,
            edgecolors="none",
        )

    ax.axhline(0.0, color="#333333", linewidth=0.8, linestyle="--")
    ax.set_xlabel("Die area (um^2)")
    ax.set_ylabel("WNS (ns)")
    ax.set_title("Area vs WNS (Pareto scatter)")
    ax.legend(frameon=False, fontsize=8, loc="best")
    fig.tight_layout()

    out_base = Path(args.out)
    out_base.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_base.with_suffix(".png"), dpi=300)
    fig.savefig(out_base.with_suffix(".pdf"))
    plt.close(fig)


if __name__ == "__main__":
    main()