"""Phase-6 exit-gate measurement: active-learning selection vs a fixed (random-order) baseline.

Task: from a pool of N candidate Cu5Au2 clusters, find the lowest EMT relaxed energy with as few
"expensive" evaluations (tight EMT relaxations) as possible.
  baseline : evaluate pool members in a random order
  AL (EI)  : 3 random seeds, then GP expected improvement on structure fingerprints (discovery/active_learning.py)
Two feature settings:
  A) fingerprint of the UNRELAXED structure (hard: features weakly related to the relaxed energy)
  B) fingerprint after a very cheap pre-relaxation (8 FIRE steps) - a multi-fidelity setting
Reports the median / mean number of evaluations to hit the pool minimum over many seeds. Results are
written to docs/benchmarks/active_learning.json and reported as measured, whichever way they come out.

    python scripts/active_learning_benchmark.py [--pool 60] [--trials 30]
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from ase.calculators.emt import EMT  # noqa: E402
from ase.optimize import BFGS, FIRE  # noqa: E402

from core.representation import MatterStructure  # noqa: E402
from discovery.active_learning import select  # noqa: E402
from discovery.generators import random_cluster  # noqa: E402


def build_pool(n: int, seed: int = 0):
    rng = np.random.default_rng(seed)
    syms = ["Cu"] * 5 + ["Au"] * 2
    pool = []
    for i in range(n):
        ms = random_cluster(list(rng.permutation(syms)), seed=int(rng.integers(2**31)))
        a = ms.to_ase(); a.calc = EMT()
        FIRE(a, logfile=None).run(fmax=0.5, steps=8)
        cheap = MatterStructure.from_ase(a)
        b = ms.to_ase(); b.calc = EMT()
        BFGS(b, logfile=None).run(fmax=0.01, steps=1500)
        pool.append({"fp_unrelaxed": ms.fingerprint(), "fp_cheap": cheap.fingerprint(), "energy": float(b.get_potential_energy())})
    return pool


def evals_to_min(order_fn, pool, rng):
    best = min(p["energy"] for p in pool)
    return order_fn(pool, rng, best)


def random_order(pool, rng, best):
    for k, i in enumerate(rng.permutation(len(pool)), 1):
        if abs(pool[i]["energy"] - best) < 1e-9:
            return k
    return len(pool)


def make_al(feature):
    def al(pool, rng, best):
        seen = list(rng.choice(len(pool), 3, replace=False))
        for i in seen:
            if abs(pool[i]["energy"] - best) < 1e-9:
                return seen.index(i) + 1
        while len(seen) < len(pool):
            rest = [i for i in range(len(pool)) if i not in seen]
            res = select([pool[i][feature] for i in rest], [pool[i][feature] for i in seen], [pool[i]["energy"] for i in seen],
                         k=1, acquisition="expected_improvement", diversity_penalty=0.0)
            nxt = rest[res["chosen"][0]]
            seen.append(nxt)
            if abs(pool[nxt]["energy"] - best) < 1e-9:
                return len(seen)
        return len(pool)
    return al


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--pool", type=int, default=60); ap.add_argument("--trials", type=int, default=30)
    a = ap.parse_args()
    t0 = time.time()
    pool = build_pool(a.pool)
    res = {}
    for name, fn in [("random_baseline", random_order), ("al_ei_unrelaxed_features", make_al("fp_unrelaxed")),
                     ("al_ei_cheap_prerelax_features", make_al("fp_cheap"))]:
        ks = [evals_to_min(fn, pool, np.random.default_rng(1000 + t)) for t in range(a.trials)]
        res[name] = {"median": float(np.median(ks)), "mean": float(np.mean(ks)), "std": float(np.std(ks)), "trials": ks}
    out = {"task": "evaluations needed to find the pool minimum (lower is better)", "pool_size": a.pool, "trials": a.trials,
           "solver": "ASE EMT, BFGS fmax=0.01", "results": {k: {kk: v[kk] for kk in ("median", "mean", "std")} for k, v in res.items()},
           "raw": {k: v["trials"] for k, v in res.items()}, "elapsed_s": round(time.time() - t0, 1),
           "expected_random_mean": (a.pool + 1) / 2}
    d = ROOT / "docs" / "benchmarks"; d.mkdir(parents=True, exist_ok=True)
    (d / "active_learning.json").write_text(json.dumps(out, indent=1))
    print(json.dumps({k: out[k] for k in ("pool_size", "trials", "results", "expected_random_mean", "elapsed_s")}, indent=1))


if __name__ == "__main__":
    main()