"""Build the leaderboard notebook's two data files.

    python build_data.py results.parquet out/*.jsonl

`results.parquet` is the leaderboard's results dataset
(proxima-fusion/constellaration-bench-results, parquet conversion). From it,
data/entries.parquet keeps one row per submission and one per MHD-stable
point: who, when, the recorded objectives, feasibility and score. From the
JSON lines of scripts/rescore.py, data/rescored.parquet keeps one row per
boundary we evaluated, with the recorded values beside ours; a boundary
evaluated again keeps one row, and the build stops when the flux-compression
files cover only part of the first- or second-ranked front. For the
evaluate notebook, ../evaluate/data/leading-front.json keeps the feasible
(L∇B, A) points of the leading MHD-stable submission. No file holds
boundary coefficients.
"""

import json
import sys
from pathlib import Path

import pandas as pd

TOLERANCE = 0.01
RUN_KINDS = {"kernel": "kernel image", "top": "leaderboard environment", "space": "leaderboard environment", "flux": "flux compression only"}


def rows_from(path: Path):
    environment, recorded = None, {}
    for line in path.read_text().splitlines():
        record = json.loads(line)
        if record["kind"] == "environment":
            environment, recorded = record, record.get("recorded", {})
            continue
        entry = record["entry"]
        user, _, submitted = entry.rpartition("@")
        known = recorded.get(entry, {})
        row = {
            "source_file": path.name,
            "run": next((label for key, label in RUN_KINDS.items() if path.name.startswith(key)), path.stem),
            "python": environment["python"],
            "booz_xform": environment["packages"]["booz-xform"],
            "user": user,
            "submitted": submitted,
            "problem": record["problem"],
            "index": record["index"],
            "n_field_periods": record.get("n_field_periods"),
            "poloidal_modes": (record.get("modes") or [None])[0],
            "seconds": record["seconds"],
            "max_rss_mib": record["max_rss_mib"],
            "error": record.get("error"),
            "feasibility": record.get("feasibility"),
            "score": record.get("score"),
            "objective": record.get("objective"),
            "flux_compression_pest": record.get("flux_compression_pest_jacobian"),
            "flux_compression_vmec": record.get("flux_compression_vmec_jacobian"),
        }
        for name, value in (record.get("metrics") or {}).items():
            row[name] = value
        if known:
            if record["problem"] == "mhd_stable":
                feasibilities = json.loads(known["feasibilities"]) if isinstance(known["feasibilities"], str) else known["feasibilities"]
                objectives = json.loads(known["objectives"]) if isinstance(known["objectives"], str) else known["objectives"]
                row["recorded_feasibility"] = feasibilities[record["index"]]
                row["recorded_gradient_scale_length"] = objectives[record["index"]][0][0]
                row["recorded_aspect_ratio"] = objectives[record["index"]][1][0]
                row["recorded_entry_score"] = known["score"]
            else:
                row["recorded_feasibility"] = known["feasibility"]
                row["recorded_objective"] = known["objective"]
                row["recorded_entry_score"] = known["score"]
        yield row


def entries_from(results: Path) -> pd.DataFrame:
    frame = pd.read_parquet(results)
    frame = frame[(frame["user"] != "test") & frame["evaluated"]]
    rows = []
    for _, entry in frame.iterrows():
        base = {"user": entry["user"], "submitted": entry["submission_time"], "problem": entry["problem_type"], "entry_score": entry["score"]}
        if entry["problem_type"] == "mhd_stable":
            objectives = json.loads(entry["objectives"]) if entry["objectives"] else []
            feasibilities = json.loads(entry["feasibilities"]) if entry["feasibilities"] else []
            for index, (objective, feasibility) in enumerate(zip(objectives, feasibilities)):
                rows.append({**base, "index": index, "gradient_scale_length": objective[0][0], "aspect_ratio": objective[1][0], "feasibility": feasibility})
        else:
            rows.append({**base, "index": 0, "objective": entry["objective"], "feasibility": entry["feasibility"]})
    return pd.DataFrame(rows)


def leading_front(entries: pd.DataFrame) -> dict:
    points = entries[entries["problem"] == "mhd_stable"]
    leader = points.loc[points["entry_score"].idxmax()]
    chosen = points[(points["user"] == leader["user"]) & (points["submitted"] == leader["submitted"])]
    feasible = chosen[chosen["feasibility"] <= TOLERANCE].sort_values("index")
    return {
        "user": leader["user"],
        "submitted": leader["submitted"],
        "recorded_score": float(leader["entry_score"]),
        "columns": ["minimum_normalized_magnetic_gradient_scale_length", "aspect_ratio"],
        "points": feasible[["gradient_scale_length", "aspect_ratio"]].astype(float).to_numpy().tolist(),
    }


def without_reruns(rescored: pd.DataFrame) -> pd.DataFrame:
    """One row per boundary and run kind; a successful row wins over a failed one, then the later file."""
    if rescored.empty:
        return rescored
    order = rescored.assign(_failed=rescored["error"].notna(), _position=range(len(rescored)))
    order = order.sort_values(["_failed", "_position"], ascending=[False, True])
    kept = order.drop_duplicates(["run", "user", "submitted", "problem", "index"], keep="last")
    return kept.sort_values("_position").drop(columns=["_failed", "_position"])


def check_flux_fronts(entries: pd.DataFrame, rescored: pd.DataFrame) -> None:
    """Fail when the flux-compression run covers only part of the first- or second-ranked front.

    Ranks 1 and 2 are evaluated point by point in shards; ranks 3 and 4 are
    samples and are not checked. scripts/rescore.py ranks after dropping
    resubmissions of identical boundaries, which have identical scores.
    """
    flux = rescored[rescored["run"] == RUN_KINDS["flux"]]
    if flux.empty:
        return
    points = entries[entries["problem"] == "mhd_stable"]
    ranked = points.groupby(["user", "submitted"], as_index=False).agg(score=("entry_score", "first"), points=("index", "size"))
    ranked = ranked.sort_values(["score", "submitted"], ascending=[False, True]).drop_duplicates("score")
    for rank, leader in enumerate(ranked.head(2).itertuples(), start=1):
        found = flux[(flux["user"] == leader.user) & (flux["submitted"] == leader.submitted)]["index"]
        missing = sorted(set(range(leader.points)) - set(found))
        if missing:
            sys.exit(f"flux compression of rank {rank} ({leader.user}) lacks {len(missing)} of {leader.points} points, from index {missing[0]}")


def main() -> None:
    data = Path(__file__).parent / "data"
    data.mkdir(exist_ok=True)
    entries = entries_from(Path(sys.argv[1]))
    entries.to_parquet(data / "entries.parquet", compression="zstd", index=False)
    print(f"{len(entries)} rows -> {data / 'entries.parquet'}")
    front = Path(__file__).parent.parent / "evaluate" / "data" / "leading-front.json"
    front.write_text(json.dumps(leading_front(entries)) + "\n")
    print(f"leading MHD-stable front -> {front}")
    rescored = without_reruns(pd.DataFrame([row for name in sys.argv[2:] for row in rows_from(Path(name))]))
    check_flux_fronts(entries, rescored)
    rescored.to_parquet(data / "rescored.parquet", compression="zstd", index=False)
    print(f"{len(rescored)} rows -> {data / 'rescored.parquet'}")


if __name__ == "__main__":
    main()