"""Refine BOTH active families at selected budget faces after the coarse scan.

All evaluated protocols stay in the global catalogue, even when infeasible for
the particular face being optimized. Envelopes subsequently enforce both caps.
Convergence is recorded honestly; these local searches do not certify optima.
Rerunning filters prior refinement rows and rebuilds from the coarse catalogue.
"""
import csv
import json
import numpy as np
from scipy.optimize import minimize
import matched_protection_budgets as m


def load_coarse():
    n = len(list(m.catalogue()))*len(m.DEADLINES)*len(m.Z_RATES)
    rows = []
    for raw in csv.DictReader((m.OUT/"protocol-catalogue.csv").open()):
        if int(raw["id"]) >= n:
            continue
        row = {k: (v if k == "kind" else float(v)) for k, v in raw.items()}
        row["id"] = int(row["id"])
        rows.append(row)
    assert len(rows) == n and [r["id"] for r in rows] == list(range(n))
    return rows


def main():
    rows = load_coarse()
    logs = []
    fields = list(m.Protocol.__dataclass_fields__)
    def key(W, kz, spec):
        return (W, kz, spec.kind)+tuple(round(float(getattr(spec, k)), 12) for k in fields[1:])
    cache = {key(r["deadline"], r["kz"], m.Protocol(**{k: r[k] for k in fields})): r for r in rows}
    faces = [(2., .02, 20., 20.), (2., .02, 80., 40.),
             (2., .02, 200., 80.), (.5, .02, 20., 20.), (2., .2, 80., 40.)]
    for W, kz, Bg, Br in faces:
        for prep in m.PREP_TIMES:
            core = W-prep-.12
            if core <= 0 or Br < m.PREP_RATE*prep:
                continue
            for kind in ["continuous", "switched"]:
                def evaluate(x):
                    if kind == "continuous":
                        spec = m.Protocol(kind, prep, float(x[0]), float(x[1]))
                    else:
                        spec = m.Protocol(kind, prep, *map(float, x))
                    kk = key(W, kz, spec)
                    if kk not in cache:
                        a = m.simulate(spec, W=W, kz=kz)
                        a["id"] = len(rows); rows.append(a); cache[kk] = a
                    return cache[kk]
                feasible = [r for r in rows if r["deadline"] == W and r["kz"] == kz and
                            r["kind"] == kind and r["prep"] == prep and
                            r["coupling_budget_used"] <= Bg+1e-9 and
                            r["aux_reset_budget_used"] <= Br+1e-9]
                feasible.sort(key=lambda r: -r["Fe"])
                if kind == "continuous":
                    gcap = min(10., np.sqrt(Bg/core))
                    rcap = min(40., (Br-m.PREP_RATE*prep)/core)
                    bounds = [(0., gcap), (0., rcap)]
                    starts = [[gcap, rcap], [.5*gcap, .5*rcap]]
                    if feasible:
                        starts.append([feasible[0]["g"], feasible[0]["gamma"]])
                    method = "L-BFGS-B"
                    constraints = ()
                else:
                    bounds = [(.5, 10.), (1., 40.), (.3, 2.5), (.1, 4.), (0., .5), (0., .95)]
                    starts = [[r[k] for k in fields[2:]] for r in feasible[:2]]
                    if not starts:
                        continue
                    method = "SLSQP"
                    constraints = [{"type": "ineq", "fun": lambda x: Bg-evaluate(x)["coupling_budget_used"]},
                                   {"type": "ineq", "fun": lambda x: Br-evaluate(x)["aux_reset_budget_used"]}]
                for start in starts:
                    before = len(rows)
                    result = minimize(lambda x: -evaluate(x)["Fe"], start, method=method,
                        bounds=bounds, constraints=constraints,
                        options={"maxiter": 60, "ftol": 1e-10})
                    end = evaluate(result.x)
                    logs.append({"deadline": W, "kz": kz, "coupling_cap": Bg, "reset_cap": Br,
                        "prep": prep, "kind": kind, "method": method,
                        "success": bool(result.success), "message": str(result.message),
                        "iterations": int(result.nit), "new_evaluations": len(rows)-before,
                        "end_id": end["id"], "end_Fe": end["Fe"],
                        "end_coupling_violation": max(0., end["coupling_budget_used"]-Bg),
                        "end_reset_violation": max(0., end["aux_reset_budget_used"]-Br)})
            print("refined face", W, kz, Bg, Br, "prep", prep, "rows", len(rows), flush=True)
    record = {"scope": __doc__, "faces": faces, "runs": logs,
              "successful_runs": sum(x["success"] for x in logs),
              "all_runs": len(logs), "rows_after_refinement": len(rows),
              "bounds_switched": {"g": [.5, 10.], "gamma": [1., 40.],
                                  "angle": [.3, 2.5], "reset_exposure": [.1, 4.],
                                  "wait": [0., .5], "offset": [0., .95]}}
    (m.OUT/"refinement.json").write_text(json.dumps(record, ensure_ascii=False, indent=2)+"\n")
    m.finalize(rows, refinement={"file": "refinement.json", "faces": faces,
                                "runs": len(logs), "successful_runs": record["successful_runs"],
                                "certified_global_optimum": False})
    print(json.dumps({k: record[k] for k in ["successful_runs", "all_runs", "rows_after_refinement"]}))


if __name__ == "__main__":
    main()