"""
08_dps_level_k_bisection.py

General, RESUMABLE bisection for the DPS level-k noise-robustness
threshold of the Tiles state family. Each run performs STEPS_PER_RUN
bisection steps and saves progress to a JSON state file, so you can call
it repeatedly (e.g. in a shell loop, or across separate sessions) without
losing progress -- useful since each SDP solve can take anywhere from
under a second (k=2) to a few minutes (k=3, with SCS) depending on k,
your hardware, and the solver.

Usage:
    python3 08_dps_level_k_bisection.py

Configure LEVEL, SOLVER, STEPS_PER_RUN, and EPS below.
The state file is named dps_level{LEVEL}_bisection_state.json.

--------------------------------------------------------------------
Progress already made in the original chat session for LEVEL=3 (8 SCS
solves, ~135-227s each) is included alongside this script as
dps_level3_bisection_state.json:

    bracket so far: [0.90982, 0.91080]   (i.e. p_c ~ 0.910-0.911)

Just run this script (with LEVEL=3, the default) to continue narrowing
it -- it will pick up automatically from that saved state. Delete the
state file to start over, or change LEVEL to try a different extension
order (4, 5, ... but see README.md for how fast the PPT-constraint size,
and hence the cost, grows: 3*3^k).
--------------------------------------------------------------------
"""
import json
import os
import time
import cvxpy as cp
from common import noisy_tiles
from dps_hierarchy import build_dps_problem, dps_feasible

LEVEL = 3
SOLVER = cp.SCS            # swap to cp.MOSEK if available -- likely much faster
STEPS_PER_RUN = 1          # raise this if your machine/solver is fast enough
EPS = 1e-5                 # solver tolerance; tighten once you have a rough bracket
DEFAULT_BRACKET = (0.70, 0.951)  # 0.951 is a proven-safe upper bound (= DPS level-2 threshold,
                                  # since DPS level 3 can only detect at <= that noise level)

STATE_FILE = f"dps_level{LEVEL}_bisection_state.json"

if os.path.exists(STATE_FILE):
    with open(STATE_FILE) as f:
        state = json.load(f)
    print(f"Resuming from saved state: bracket [{state['lo']:.5f}, {state['hi']:.5f}], "
          f"{state['iter']} iterations so far.")
else:
    state = {"lo": DEFAULT_BRACKET[0], "hi": DEFAULT_BRACKET[1], "iter": 0, "log": []}
    print(f"Starting fresh: bracket {DEFAULT_BRACKET}")

prob, rho_param, sigma = build_dps_problem(d=3, k=LEVEL)
print(f"sigma shape: {sigma.shape},  PPT-constraint (PSD cone) size: "
      f"{3 * 3 ** LEVEL} x {3 * 3 ** LEVEL}\n")

for _ in range(STEPS_PER_RUN):
    lo, hi = state["lo"], state["hi"]
    mid = (lo + hi) / 2
    t0 = time.time()
    feasible = dps_feasible(prob, rho_param, noisy_tiles(mid), solver=SOLVER,
                             eps=EPS, max_iters=20000, warm_start=True)
    dt = time.time() - t0
    if feasible:
        state["lo"] = mid
    else:
        state["hi"] = mid
    state["iter"] += 1
    state["log"].append({"iter": state["iter"], "p": mid, "feasible": feasible,
                          "time_s": round(dt, 1), "status": prob.status})
    print(f"[iter {state['iter']}] p={mid:.5f}  feasible={feasible}  "
          f"status={prob.status}  ({dt:.1f}s)   "
          f"bracket now [{state['lo']:.5f}, {state['hi']:.5f}]")

with open(STATE_FILE, "w") as f:
    json.dump(state, f, indent=2)

print(f"\nCurrent bracket: [{state['lo']:.5f}, {state['hi']:.5f}]  "
      f"(width {state['hi'] - state['lo']:.5f})")
print("Run again to continue narrowing it further (progress is saved).")