#!/usr/bin/env python
import argparse
import json
import math
import statistics
from pathlib import Path


METHODS = ("sft", "random", "greedy_verifier", "oracle_grid")
METRICS = ("hypervolume", "mip", "preference_regret", "near_optimal", "unique")
DIRECTIONS = {"up_up": (1, 1), "up_down": (1, -1),
              "down_up": (-1, 1), "down_down": (-1, -1)}


def ci(xs):
    sd = statistics.stdev(xs)
    return {"mean": statistics.fmean(xs), "sd": sd, "ci95": 2.776 * sd / math.sqrt(len(xs))}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", type=Path, required=True)
    ap.add_argument("--output", type=Path, required=True)
    ap.add_argument("--markdown", type=Path, required=True)
    args = ap.parse_args()
    runs = [json.loads(p.read_text()) for p in sorted(args.input.glob("seed_*.json"))]
    if len(runs) != 5 or len({r["split_hash"] for r in runs}) != 1:
        raise SystemExit("expected five seeds with one split")
    summary = {m: {k: ci([r["results"][m][k] for r in runs]) for k in METRICS}
               for m in METHODS}
    paired, relative = {}, {}
    for base in ("random", "greedy_verifier"):
        paired[base] = {k: ci([r["results"]["sft"][k] - r["results"][base][k]
                               for r in runs]) for k in METRICS}
        relative[base] = {k: ci([100 * (r["results"]["sft"][k] - r["results"][base][k])
                                 / r["results"][base][k] for r in runs])
                          for k in ("hypervolume", "mip")}
    by_direction = {}
    for name, direction in DIRECTIONS.items():
        by_direction[name] = {m: {metric: ci([
            statistics.fmean(q[metric] for q in r["results"][m]["tasks"]
                             if tuple(q["directions"]) == direction) for r in runs])
            for metric in ("hypervolume", "mip")} for m in METHODS}
        by_direction[name]["sft_minus_greedy"] = {metric: ci([
            statistics.fmean(q[metric] for q in r["results"]["sft"]["tasks"]
                             if tuple(q["directions"]) == direction)
            - statistics.fmean(q[metric] for q in r["results"]["greedy_verifier"]["tasks"]
                               if tuple(q["directions"]) == direction) for r in runs])
            for metric in ("hypervolume", "mip")}
    verifier_rmse = {p: ci([r["val_rmse"][p] for r in runs]) for p in ("Egc", "Egb")}
    output = {"seeds": [r["seed"] for r in runs], "split_hash": runs[0]["split_hash"],
              "data_sha256": runs[0]["data_sha256"], "n_eval_tasks": runs[0]["n_eval_tasks"],
              "n_preference_decisions": runs[0]["n_preference_decisions"],
              "verifier_rmse": verifier_rmse, "summary": summary, "by_direction": by_direction,
              "sft_minus": paired, "sft_relative_percent": relative}
    args.output.write_text(json.dumps(output, indent=2) + "\n")
    lines = ["# PolyEdit multi-objective Hypervolume and MIP", "",
             "Values are five-seed means with 95% t confidence intervals.", "",
             "| Method | Hypervolume | MIP | Preference regret | Near-optimal | Unique |",
             "|---|---:|---:|---:|---:|---:|"]
    for method in METHODS:
        row = summary[method]
        cell = lambda k: f"{row[k]['mean']:.3f} ± {row[k]['ci95']:.3f}"
        lines.append(f"| {method} | {cell('hypervolume')} | {cell('mip')} | "
                     f"{cell('preference_regret')} | {cell('near_optimal')} | {cell('unique')} |")
    lines += ["", "## Direction breakdown", "",
              "| Direction | Method | Hypervolume | MIP |", "|---|---|---:|---:|"]
    for name in DIRECTIONS:
        for method in METHODS:
            row = by_direction[name][method]
            lines.append(f"| {name} | {method} | {row['hypervolume']['mean']:.3f} ± "
                         f"{row['hypervolume']['ci95']:.3f} | {row['mip']['mean']:.3f} ± "
                         f"{row['mip']['ci95']:.3f} |")
    args.markdown.write_text("\n".join(lines) + "\n")


if __name__ == "__main__":
    main()