"""Score ConStellaration leaderboard submissions again with the installed evaluator.

Each boundary goes through the forward model once, with the settings the
official evaluator uses for its problem. The script records every metric, the
objective, the normalised constraint violations and the time and memory each
evaluation took, one JSON object per line. Scores of a whole submission are
computed from these lines afterwards by `leaderboard/build_data.py` and the
leaderboard notebook.

    python rescore.py --entries baselines --out baselines.jsonl
    python rescore.py --entries top --problem mhd_stable --top 4 --jobs 8 --out p3.jsonl

The leaderboard results are read from Hugging Face at run time and are not
copied anywhere.
"""

from __future__ import annotations

import argparse
import concurrent.futures
import importlib.metadata
import io
import json
import platform
import resource
import sys
import time
import traceback
import urllib.request

RESULTS_URL = (
    "https://huggingface.co/datasets/proxima-fusion/constellaration-bench-results"
    "/resolve/refs%2Fconvert%2Fparquet/default/train/0000.parquet"
)
REFERENCE_BOUNDARY_URL = (
    "https://raw.githubusercontent.com/proximafusion/constellaration/v0.2.6"
    "/hugging_face_competition/inputs/boundary.json"
)
PROBLEMS = ("geometrical", "simple_to_build", "mhd_stable")


def problem_object(name: str):
    from constellaration import problems

    return {
        "geometrical": problems.GeometricalProblem,
        "simple_to_build": problems.SimpleToBuildQIStellarator,
        "mhd_stable": problems.MHDStableQIStellarator,
    }[name]()


def settings_for(name: str):
    from constellaration import forward_model

    settings = forward_model.ConstellarationSettings
    return settings.default_high_fidelity_skip_qi() if name == "geometrical" else settings.default_high_fidelity()


def flux_compression_both_weights(equilibrium, settings) -> tuple[float, float]:
    """Flux compression in bad curvature averaged with the VMEC and the PEST Jacobian.

    The first value is what constellaration 0.2.6 computes. The second changes
    only the averaging weight, as constellaration 0.3.0 does (PR #109), and
    leaves everything else, including the equilibrium, as it is.
    """
    import numpy as np
    from constellaration.mhd import turbulent_transport, vmec_utils
    from simsopt import mhd

    vmec = vmec_utils.as_simsopt_vmec(equilibrium)
    alpha, theta = turbulent_transport._set_up_surface_grid(
        settings=settings, n_field_periods=equilibrium.nfp, is_stellarator_symmetric=not equilibrium.lasym
    )
    data = mhd.vmec_fieldlines(vs=vmec, s=settings.normalized_toroidal_flux, alpha=alpha, theta1d=theta)
    grad_r_dot_grad_r = equilibrium.Aminor_p**2 / (4 * settings.normalized_toroidal_flux[:, None, None]) * data.grad_s_dot_grad_s
    curvature_drift = data.B_cross_kappa_dot_grad_alpha * np.sign(equilibrium.phipf[-1])
    compression = np.heaviside(curvature_drift, 0) * grad_r_dot_grad_r

    def averaged(jacobian):
        weights = np.abs(jacobian) / np.abs(jacobian).mean(axis=(1, 2))[:, None, None]
        return float(np.mean(compression * weights, axis=(1, 2))[0])

    return averaged(data.sqrt_g_vmec), averaged(data.edge_toroidal_flux_over_2pi / data.B_sup_theta_pest)


def evaluate_boundary(task: dict) -> dict:
    """One forward-model run in a fresh worker process."""
    from constellaration import forward_model
    from constellaration.geometry import surface_rz_fourier

    started = time.perf_counter()
    record = {key: task[key] for key in ("entry", "problem", "index")}
    try:
        boundary = surface_rz_fourier.SurfaceRZFourier.model_validate_json(task["boundary"])
        record["n_field_periods"] = boundary.n_field_periods
        record["modes"] = list(boundary.r_cos.shape)
        problem = problem_object(task["problem"])
        if task.get("flux_compression_only"):
            # Without the Boozer transform and QI residual, which the
            # correction does not touch and which need about 3 GiB.
            settings = settings_for(task["problem"]).model_copy(update={"boozer_preset_settings": None, "qi_settings": None})
            metrics, equilibrium = forward_model.forward_model(boundary, settings=settings)
            record["metrics"] = metrics.model_dump(mode="json")
            vmec_weight, pest_weight = flux_compression_both_weights(equilibrium, settings.turbulent_settings)
            record["flux_compression_vmec_jacobian"] = vmec_weight
            record["flux_compression_pest_jacobian"] = pest_weight
            record["objectives"] = problem.get_objectives(metrics)
        else:
            metrics, _ = forward_model.forward_model(boundary, settings=settings_for(task["problem"]))
            record["metrics"] = metrics.model_dump(mode="json")
            record["violations"] = problem._normalized_constraint_violations(metrics).tolist()
            record["feasibility"] = problem.compute_feasibility(metrics)
            record["feasible"] = problem.is_feasible(metrics)
            if task["problem"] == "mhd_stable":
                record["objectives"] = problem.get_objectives(metrics)
            else:
                record["objective"] = problem.get_objective(metrics)[0]
                record["score"] = problem._score(metrics) if record["feasible"] else 0.0
    except Exception as error:  # VMEC++ non-convergence and QI failures are results too.
        record["error"] = f"{type(error).__name__}: {error}"
        record["traceback"] = traceback.format_exc(limit=3)
    record["seconds"] = round(time.perf_counter() - started, 2)
    record["max_rss_mib"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss // (1024 if sys.platform != "darwin" else 1024 * 1024)
    return record


def load_results():
    import pandas as pd

    with urllib.request.urlopen(RESULTS_URL, timeout=120) as response:
        frame = pd.read_parquet(io.BytesIO(response.read()))
    return frame[frame["user"] != "test"]


def select_entries(frame, which: str, problem: str | None, top: int):
    if which == "baselines":
        # The first submission per problem by the paper's first author holds the
        # ALM-NGOpt baselines of Cadena et al. (2025), Tables 3 and 5.
        chosen = frame[frame["user"] == "scadena-pf"].sort_values("submission_time").groupby("problem_type").head(1)
    else:
        candidates = frame[frame["problem_type"].isin([problem] if problem else PROBLEMS)]
        candidates = candidates.drop_duplicates(subset=["problem_type", "boundary_json"])
        chosen = candidates.sort_values("score", ascending=False).groupby("problem_type").head(top)
    return chosen


def tasks_for(entries) -> list[dict]:
    tasks = []
    for _, row in entries.iterrows():
        entry = f"{row['user']}@{row['submission_time']}"
        if row["problem_type"] == "mhd_stable":
            boundaries = json.loads(row["boundary_json"])
        else:
            boundaries = [row["boundary_json"]]
        for index, boundary in enumerate(boundaries):
            tasks.append({"entry": entry, "problem": row["problem_type"], "index": index, "boundary": boundary})
    return tasks


def environment() -> dict:
    packages = ("constellaration", "vmecpp", "booz-xform", "simsopt", "numpy", "scipy", "jax", "pymoo")
    versions = {}
    for name in packages:
        try:
            versions[name] = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            versions[name] = None
    return {"python": platform.python_version(), "machine": platform.machine(), "system": platform.system(), "packages": versions}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--entries", choices=("baselines", "top", "reference"), default="baselines")
    parser.add_argument("--problem", choices=PROBLEMS)
    parser.add_argument("--top", type=int, default=1)
    parser.add_argument("--rank", type=int, help="with --entries top: only the entry at this rank (1 = best)")
    parser.add_argument("--jobs", type=int, default=1)
    parser.add_argument("--flux-compression-only", action="store_true", help="MHD-stable only: skip QI and compute the flux compression with both Jacobian weights")
    parser.add_argument("--limit", type=int, help="evaluate only the first N boundaries")
    parser.add_argument("--sample", type=int, help="evaluate N boundaries drawn at random (seed 0)")
    parser.add_argument("--shard", help="I/N: evaluate every N-th boundary starting at I (0-based)")
    parser.add_argument("--out", required=True)
    args = parser.parse_args()

    if args.entries == "reference":
        with urllib.request.urlopen(REFERENCE_BOUNDARY_URL, timeout=60) as response:
            boundary = response.read().decode()
        tasks = [{"entry": "reference", "problem": p, "index": 0, "boundary": boundary} for p in ("geometrical", "simple_to_build")]
        recorded = {}
    else:
        entries = select_entries(load_results(), args.entries, args.problem, max(args.top, args.rank or 0))
        if args.rank:
            entries = entries.sort_values("score", ascending=False).iloc[[args.rank - 1]]
        tasks = tasks_for(entries)
        recorded = {
            f"{row['user']}@{row['submission_time']}": {
                "problem": row["problem_type"],
                "score": row["score"],
                "objectives": row["objectives"],
                "feasibilities": row["feasibilities"],
                "objective": row["objective"],
                "feasibility": row["feasibility"],
            }
            for _, row in entries.iterrows()
        }

    for task in tasks:
        task["flux_compression_only"] = args.flux_compression_only
    if args.sample and args.sample < len(tasks):
        import random

        tasks = sorted(random.Random(0).sample(tasks, args.sample), key=lambda t: t["index"])
    if args.shard:
        start, step = (int(x) for x in args.shard.split("/"))
        tasks = tasks[start::step]
    if args.limit:
        tasks = tasks[: args.limit]

    with open(args.out, "w") as out:
        out.write(json.dumps({"kind": "environment", **environment(), "recorded": recorded}) + "\n")
        out.flush()
        # max_tasks_per_child=1 (Python 3.11+) gives each evaluation its own
        # process, so the memory figure belongs to one evaluation.
        options = {"max_tasks_per_child": 1} if sys.version_info >= (3, 11) else {}
        with concurrent.futures.ProcessPoolExecutor(max_workers=args.jobs, **options) as pool:
            for record in pool.map(evaluate_boundary, tasks):
                out.write(json.dumps({"kind": "boundary", **record}) + "\n")
                out.flush()
                print(record["entry"], record["problem"], record["index"], record.get("feasibility"), record.get("error", ""), record["seconds"], "s", flush=True)


if __name__ == "__main__":
    main()