"""Oracle for the multi-instance SOS task: certificate + witness for each p_k."""
import json, os, sys, time
import numpy as np
from scipy.optimize import minimize
from sos_common import (M, N_VARS, P_MONOMIALS, is_positive_definite, eval_poly_exact)
import reference_solver as R

HERE = os.path.dirname(os.path.abspath(__file__))
INST = os.environ.get("SOS_INSTANCE_DIR", os.path.join(HERE, "..", "instance"))
OUT  = os.environ.get("SOS_OUTPUT_DIR",  os.path.join(HERE, "..", "reference_outputs"))

def solve_one(p, log):
    A, ai, _ = R.build_linear_map()
    S, AS, null, rank = R.sym_basis_nullspace(A)
    rp = np.zeros(len(P_MONOMIALS))
    for a, c in p.items(): rp[ai[a]] = c
    r1 = np.zeros(len(P_MONOMIALS)); r1[ai[(0,)*N_VARS]] = 1.0
    Qp, Q1 = R.solve_particular(AS, S, rp), R.solve_particular(AS, S, r1)
    B = R.mats_from_sym(S, null)
    lo, hi, tw = -400.0, 200.0, None
    t, l = R.max_lambda_min(Qp - lo*Q1, B)
    if l <= 0: raise RuntimeError("no feasible start")
    tw = t
    for _ in range(45):
        mid = 0.5*(lo+hi)
        tm, lm = R.max_lambda_min(Qp - mid*Q1, B, tw)
        if lm > 0: lo, tw = mid, tm
        else: hi = mid
        if hi-lo < 1e-9: break
    gstar = lo
    for back, db in ((3e-6,19),(1e-5,19),(3e-5,19),(1e-4,19),(1e-3,19),(1e-2,18)):
        g = gstar - back; den = 1 << db; G = int(round(g*den))
        Q0 = Qp - (G/den)*Q1
        t, _ = R.max_lambda_min(Q0, B, tw)
        try: t, _ = R.recenter(Q0, B, t)
        except Exception: pass
        Qf = Q0 + np.tensordot(t, B, axes=(0,0))
        Qint, dr, Gr = R.rationalise(Qf, den, p, G)
        if is_positive_definite(Qint):
            return gstar, {"den": dr, "G": Gr, "Q": Qint}
    raise RuntimeError("no exact PD certificate")

def poly_grad_hess(p):
    """Analytic gradient and Hessian evaluators for the coefficient dict."""
    terms = [(np.array(e), float(c)) for e, c in p.items()]
    def val(x):
        return sum(c * np.prod(x ** e) for e, c in terms)
    def grad(x):
        g = np.zeros(N_VARS)
        for e, c in terms:
            for k in range(N_VARS):
                if e[k]:
                    ek = e.copy(); ek[k] -= 1
                    g[k] += c * e[k] * np.prod(x ** ek)
        return g
    def hess(x):
        H = np.zeros((N_VARS, N_VARS))
        for e, c in terms:
            for k in range(N_VARS):
                if not e[k]:
                    continue
                for l in range(N_VARS):
                    if l == k:
                        if e[k] >= 2:
                            ek = e.copy(); ek[k] -= 2
                            H[k, k] += c * e[k] * (e[k] - 1) * np.prod(x ** ek)
                    elif e[l]:
                        ek = e.copy(); ek[k] -= 1; ek[l] -= 1
                        H[k, l] += c * e[k] * e[l] * np.prod(x ** ek)
        return H
    return val, grad, hess


def witness(p, G, den, log, gap_tol_num=1, gap_tol_den=100000, scale=1.0):
    val, grad, hess = poly_grad_hess(p)
    best, bx = np.inf, None
    rng = np.random.default_rng(7)
    for _ in range(80):
        r = minimize(val, rng.uniform(-1.6, 1.6, N_VARS),
                     method="Nelder-Mead",
                     options={"xatol": 1e-8, "fatol": 1e-8, "maxiter": 3000})
        x = r.x
        for _ in range(30):            # Newton polish to machine precision
            g = grad(x)
            H = hess(x)
            try:
                step = np.linalg.solve(H, g)
            except np.linalg.LinAlgError:
                break
            x = x - step
            if np.linalg.norm(step) < 1e-14:
                break
        f = val(x)
        if f < best and np.linalg.norm(grad(x)) < 1e-8:
            best, bx = float(f), x
    from fractions import Fraction
    tol = Fraction(gap_tol_num, gap_tol_den) * Fraction(scale).limit_denominator(10**6)
    for D in (10**6, 10**7, 10**8):
        num = [int(round(v * D)) for v in bx]
        gap = eval_poly_exact(p, num, [D] * N_VARS) - Fraction(G, den)
        if 0 <= gap <= tol:
            return num, [D] * N_VARS, float(gap), best
    return num, [D] * N_VARS, float(gap), best

def main():
    t0 = time.time()
    def log(m): print(f"[{time.time()-t0:6.1f}s] {m}", flush=True)
    data = json.load(open(os.path.join(INST, "polynomials.json")))
    sub, report = {}, {}
    for key, terms in data["polynomials"].items():
        p = {tuple(t["exponents"]): int(t["coeff"]) for t in terms}
        gstar, cert = solve_one(p, log)
        scale = max(1.0, abs(cert["G"]) / cert["den"] / 10.0)
        num, wden, gap, emin = witness(p, cert["G"], cert["den"], log, scale=scale)
        cert["x0_num"], cert["x0_den"] = num, wden
        sub[key] = cert
        g = cert["G"]/cert["den"]
        report[key] = {"gamma_star": gstar, "gamma_certified": g,
                       "empirical_min": emin, "enclosure_width": gap,
                       "den": cert["den"]}
        log(f"{key}: gamma*={gstar:.10f} certified={g:.10f} "
            f"emin={emin:.10f} width={gap:.2e}")
    os.makedirs(OUT, exist_ok=True)
    json.dump(sub, open(os.path.join(OUT, "solution_multi.json"), "w"))
    report["wall_seconds"] = time.time()-t0
    json.dump(report, open(os.path.join(OUT, "solver_report_multi.json"), "w"), indent=2)
    log("done")

if __name__ == "__main__":
    main()