#!/usr/bin/env python3
"""
Aggregate validation results across all relations into a consolidated table.

Searches for validation_results.json files in the output directory structure:
    step4_outputs/{relation}/run_*/validation_results.json

Usage:
    python -m experiment.scripts.aggregate_results
    python -m experiment.scripts.aggregate_results --output_base ./step4_outputs
    python -m experiment.scripts.aggregate_results --format csv --out results.csv
"""

import argparse
import json
import math
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from experiment.config.relation_config import get_relation_config, list_relation_keys


def find_latest_results(output_base: str, relation: str) -> dict | None:
    """Find the most recent validation_results.json for a relation."""
    rel_dir = os.path.join(output_base, relation)
    if not os.path.isdir(rel_dir):
        return None

    # Find all run dirs, sort by name (timestamp-based)
    run_dirs = sorted(
        [d for d in os.listdir(rel_dir) if d.startswith("run_")],
        reverse=True,
    )

    for run_dir in run_dirs:
        results_path = os.path.join(rel_dir, run_dir, "validation_results.json")
        if os.path.exists(results_path):
            with open(results_path) as f:
                data = json.load(f)
            data["_run_dir"] = run_dir
            data["_path"] = results_path
            return data

    # Also check directly in rel_dir (flat structure)
    results_path = os.path.join(rel_dir, "validation_results.json")
    if os.path.exists(results_path):
        with open(results_path) as f:
            data = json.load(f)
        data["_run_dir"] = "flat"
        data["_path"] = results_path
        return data

    return None


def fmt(v, pct=False):
    """Format a metric value."""
    if v is None or (isinstance(v, float) and math.isnan(v)):
        return "N/A"
    if pct:
        return f"{v:.1%}"
    return f"{v:.3f}"


def print_markdown_table(all_results: dict):
    """Print a consolidated markdown table."""
    relations = sorted(all_results.keys())

    print("\n## Efficacy & Locality (Keyword Mention Rate)")
    print("")
    print("| Relation | Efficacy↑ | Loc_pos↑ | Loc_unrel↑ | Gen_seen↑ | Gen_unseen↑ |")
    print("|----------|-----------|----------|------------|-----------|-------------|")

    for rel in relations:
        data = all_results[rel]
        kme = data.get("kme_metrics", {})
        rc = get_relation_config(rel)

        eff = kme.get("efficacy_keyword")
        gen_seen = kme.get("generality_seen")
        gen_unseen = kme.get("generality_unseen")

        # Loc_pos: mention rate on scene_with_object (should stay high)
        summary = data.get("summary", {})
        scene_with = summary.get(rc.scene_with_object, {})
        loc_pos_ft = scene_with.get("finetuned", {}).get("mention_rate_keyword")

        # Loc_unrel: exact_match on unrelated
        loc_unrel = kme.get(f"locality/unrelated/exact_match")

        print(f"| {rel:<20} | {fmt(eff, True):>9} | {fmt(loc_pos_ft, True):>8} | "
              f"{fmt(loc_unrel):>10} | {fmt(gen_seen, True):>9} | {fmt(gen_unseen, True):>11} |")

    print("")
    print("## Caption Quality (Finetuned)")
    print("")
    print("| Relation | Efficacy Cat | With-Object Cat | Unrelated |")
    print("|----------|-------------|-----------------|-----------|")

    for rel in relations:
        data = all_results[rel]
        rc = get_relation_config(rel)
        summary = data.get("summary", {})

        eff_q = summary.get(rc.scene_no_object, {}).get("finetuned", {}).get("avg_caption_quality")
        pos_q = summary.get(rc.scene_with_object, {}).get("finetuned", {}).get("avg_caption_quality")
        unr_q = summary.get("unrelated", {}).get("finetuned", {}).get("avg_caption_quality")

        print(f"| {rel:<20} | {fmt(eff_q):>11} | {fmt(pos_q):>15} | {fmt(unr_q):>9} |")


def print_csv(all_results: dict, out_path: str = None):
    """Write results as CSV."""
    import csv
    import io

    relations = sorted(all_results.keys())
    fieldnames = [
        "relation", "efficacy_keyword", "generality_seen", "generality_unseen",
        "loc_pos_mention_rate", "loc_unrel_exact_match",
        "consistency_rouge_l", "consistency_bert_score_f1",
        "caption_quality_efficacy", "caption_quality_with_object", "caption_quality_unrelated",
        "run_dir",
    ]

    rows = []
    for rel in relations:
        data = all_results[rel]
        kme = data.get("kme_metrics", {})
        rc = get_relation_config(rel)
        summary = data.get("summary", {})

        scene_with = summary.get(rc.scene_with_object, {})
        loc_pos_ft = scene_with.get("finetuned", {}).get("mention_rate_keyword")

        rows.append({
            "relation": rel,
            "efficacy_keyword": kme.get("efficacy_keyword"),
            "generality_seen": kme.get("generality_seen"),
            "generality_unseen": kme.get("generality_unseen"),
            "loc_pos_mention_rate": loc_pos_ft,
            "loc_unrel_exact_match": kme.get("locality/unrelated/exact_match"),
            "consistency_rouge_l": kme.get("consistency/rouge_l"),
            "consistency_bert_score_f1": kme.get("consistency/bert_score_f1"),
            "caption_quality_efficacy": summary.get(rc.scene_no_object, {}).get("finetuned", {}).get("avg_caption_quality"),
            "caption_quality_with_object": summary.get(rc.scene_with_object, {}).get("finetuned", {}).get("avg_caption_quality"),
            "caption_quality_unrelated": summary.get("unrelated", {}).get("finetuned", {}).get("avg_caption_quality"),
            "run_dir": data.get("_run_dir", ""),
        })

    if out_path:
        with open(out_path, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(rows)
        print(f"CSV saved to {out_path}")
    else:
        buf = io.StringIO()
        writer = csv.DictWriter(buf, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)
        print(buf.getvalue())


def main():
    parser = argparse.ArgumentParser(description="Aggregate validation results across relations")
    parser.add_argument("--output_base", type=str, default="./step4_outputs",
                        help="Base directory containing per-relation result dirs")
    parser.add_argument("--format", type=str, choices=["markdown", "csv"], default="markdown")
    parser.add_argument("--out", type=str, default=None,
                        help="Output file path (for CSV format)")
    parser.add_argument("--relations", type=str, nargs="*", default=None,
                        help="Specific relations to include (default: all found)")
    args = parser.parse_args()

    relations = args.relations or list_relation_keys()

    all_results = {}
    for rel in relations:
        data = find_latest_results(args.output_base, rel)
        if data is not None:
            all_results[rel] = data
            print(f"  Found: {rel} ({data['_run_dir']})")
        else:
            print(f"  Missing: {rel}")

    if not all_results:
        print("\nNo results found. Run training + evaluation first.")
        return

    print(f"\nFound results for {len(all_results)}/{len(relations)} relations\n")

    if args.format == "csv":
        print_csv(all_results, args.out)
    else:
        print_markdown_table(all_results)


if __name__ == "__main__":
    main()