"""
Aggregate every band-combination run into one comparison table.

Reads, for each all_data/checkpoint*/ :
    training_history.csv  -> best val_dice_coef / val_iou_m (+ epoch)
    run_config.json       -> band list, n_bands           (written by train.py)
and the matching all_data/result*/metrics.json -> test metrics (written by evaluate.py).

    python results_table.py

Prints a table sorted by test_dice and writes all_data/runs_comparison.csv.
"""

import os
import glob
import json
import pandas as pd

ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "all_data")


def _best_val(hist_csv):
    df = pd.read_csv(hist_csv)
    if "val_dice_coef" not in df.columns or len(df) == 0:
        return None
    r = df.loc[df["val_dice_coef"].idxmax()]
    return (round(float(r["val_dice_coef"]), 4),
            round(float(r["val_iou_m"]), 4) if "val_iou_m" in df.columns else None,
            int(r["epoch"]))


def main():
    rows = []
    for cdir in sorted(glob.glob(os.path.join(ROOT, "checkpoint*"))):
        base = os.path.basename(cdir)
        name = "all" if base == "checkpoint" else base[len("checkpoint_"):]
        hist = os.path.join(cdir, "training_history.csv")
        if not os.path.exists(hist):
            continue
        bv = _best_val(hist)

        rc = os.path.join(cdir, "run_config.json")
        cfg = json.load(open(rc)) if os.path.exists(rc) else {}
        bands = cfg.get("bands")
        n_bands = cfg.get("n_bands", len(bands) if bands else "?")

        rdir = os.path.join(ROOT, "result" if name == "all" else f"result_{name}")
        mp = os.path.join(rdir, "metrics.json")
        m = json.load(open(mp)) if os.path.exists(mp) else {}

        rows.append({
            "run_name":  name,
            "n_bands":   n_bands,
            "val_dice":  bv[0] if bv else None,
            "val_iou":   bv[1] if bv else None,
            "best_ep":   bv[2] if bv else None,
            "test_dice": m.get("test_dice"),
            "test_iou":  m.get("test_iou"),
            "precision": m.get("test_precision"),
            "recall":    m.get("test_recall"),
            "img_mIoU":  m.get("per_image_mean_iou"),
            "bands":     ", ".join(bands) if bands else "(all 35)",
        })

    if not rows:
        print("No runs found under all_data/checkpoint*/  (train something first).")
        return

    df = pd.DataFrame(rows).sort_values(
        "test_dice", ascending=False, na_position="last").reset_index(drop=True)

    pd.set_option("display.width", 240)
    pd.set_option("display.max_columns", 30)
    print(df.drop(columns=["bands"]).to_string(index=False))

    out = os.path.join(ROOT, "runs_comparison.csv")
    df.to_csv(out, index=False)
    print(f"\nFull table (with band lists) -> {out}")


if __name__ == "__main__":
    main()