"""Reference (oracle) solver: exact rational SOS certificate.

Pipeline (numpy + scipy only -- there is no SDP solver in scipy, so the
semidefinite step is built here from scratch):

  1. Affine parametrisation. The identity z^T Q z = p - gamma is a linear
     system on the symmetric matrix Q. Write Q(gamma, t) = Q_p - gamma*Q_1 +
     sum_k t_k B_k, where Q_p, Q_1 are particular solutions and {B_k} is a
     basis of the null space (the classical SOS "syzygy" freedom).
  2. Semidefinite feasibility by hand. For fixed gamma, maximise a smooth
     surrogate of lambda_min(Q(gamma, t)) over t (softmin of the eigenvalues,
     exact gradient from the eigenvectors) with L-BFGS. gamma is then pushed
     up by bisection on "is max lambda_min > 0".
  3. RETREAT. At gamma = gamma* the optimal Q is singular, so no rounding of
     it can ever be positive definite. The oracle deliberately backs off to
     gamma_target < gamma* and re-centres (maximising log det) to obtain a
     Q with a comfortable spectral margin.
  4. Exact rationalisation. Scale by a power of two and round to integers,
     then repair the identity EXACTLY: for every monomial alpha a distinct
     representative entry of Q is adjusted by an integer, which changes only
     that one monomial equation. The result satisfies z^T Q z = p - gamma
     exactly, and stays positive definite provided the margin from step 3
     exceeds the rounding perturbation.
  5. Exact verification (Bareiss/Sylvester) before writing the answer.
"""

import json
import os
import time

import numpy as np
from scipy.optimize import minimize

from sos_common import (GAMMA_PARTIAL, GAMMA_TARGET, M, N_VARS, P_MONOMIALS, Z_BASIS, grade,
                        is_positive_definite, verify_submission)

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 load_p():
    with open(os.path.join(INST, "polynomial.json")) as fh:
        d = json.load(fh)
    return {tuple(t["exponents"]): int(t["coeff"]) for t in d["terms"]}


def build_linear_map():
    """A: sym(M) -> R^{n_alpha}, A(Q)_alpha = sum_{ordered i,j} Q_ij."""
    alpha_index = {a: k for k, a in enumerate(P_MONOMIALS)}
    n_alpha = len(P_MONOMIALS)
    pairs = []                      # (alpha_idx, i, j)
    for i, ei in enumerate(Z_BASIS):
        for j, ej in enumerate(Z_BASIS):
            a = tuple(ei[v] + ej[v] for v in range(N_VARS))
            pairs.append((alpha_index[a], i, j))
    A = np.zeros((n_alpha, M * M))
    for k, i, j in pairs:
        A[k, i * M + j] = 1.0
    return A, alpha_index, pairs


def sym_basis_nullspace(A):
    """Null space of A restricted to SYMMETRIC matrices, as M x M blocks."""
    # symmetric parametrisation: vector of upper-triangular entries
    iu = np.triu_indices(M)
    n_sym = iu[0].size
    S = np.zeros((M * M, n_sym))
    for c, (i, j) in enumerate(zip(*iu)):
        S[i * M + j, c] = 1.0
        S[j * M + i, c] = 1.0
    AS = A @ S                                  # (n_alpha, n_sym)
    u, s, vt = np.linalg.svd(AS)
    tol = max(AS.shape) * np.finfo(float).eps * (s[0] if s.size else 1.0)
    rank = int((s > tol).sum())
    null = vt[rank:].T                          # (n_sym, n_null)
    return S, AS, null, rank


def mats_from_sym(S, vecs):
    """Turn columns of a sym-parameter matrix into M x M matrices."""
    out = []
    for c in range(vecs.shape[1]):
        out.append((S @ vecs[:, c]).reshape(M, M))
    return np.array(out)


def solve_particular(AS, S, rhs):
    v, *_ = np.linalg.lstsq(AS, rhs, rcond=None)
    return (S @ v).reshape(M, M)


def softmin_neg(t, Q0, B, beta):
    """-softmin of eigenvalues (to be MINIMISED) plus gradient."""
    Q = Q0 + np.tensordot(t, B, axes=(0, 0))
    w, V = np.linalg.eigh(Q)
    wmin = w.min()
    ex = np.exp(-beta * (w - wmin))
    sm = wmin - np.log(ex.sum()) / beta
    p = ex / ex.sum()
    Gq = (V * p) @ V.T
    grad = np.tensordot(B, Gq, axes=([1, 2], [0, 1]))
    return -sm, -grad


def max_lambda_min(Q0, B, t0=None, beta=200.0, iters=400):
    t0 = np.zeros(B.shape[0]) if t0 is None else t0
    res = minimize(softmin_neg, t0, args=(Q0, B, beta), jac=True,
                   method="L-BFGS-B",
                   options={"maxiter": iters, "ftol": 1e-16, "gtol": 1e-12})
    Q = Q0 + np.tensordot(res.x, B, axes=(0, 0))
    return res.x, float(np.linalg.eigvalsh(Q).min())


def neg_logdet(t, Q0, B):
    Q = Q0 + np.tensordot(t, B, axes=(0, 0))
    w, V = np.linalg.eigh(Q)
    if w.min() <= 0:
        return 1e12, np.zeros(B.shape[0])
    Qi = (V / w) @ V.T
    return -float(np.log(w).sum()), -np.tensordot(B, Qi, axes=([1, 2], [0, 1]))


def recenter(Q0, B, t0, iters=300):
    res = minimize(neg_logdet, t0, args=(Q0, B), jac=True, method="L-BFGS-B",
                   options={"maxiter": iters, "ftol": 1e-16})
    Q = Q0 + np.tensordot(res.x, B, axes=(0, 0))
    return res.x, float(np.linalg.eigvalsh(Q).min())


def rationalise(Qf, den, p_coeffs, G):
    """Round to Qint/den and repair the identity exactly with integers.

    Returns (Qint, den, G). Everything is scaled by 2 after rounding: an
    off-diagonal representative (i, j) moves its monomial coefficient by
    2*delta (it appears as both Q_ij and Q_ji), so the residual has to be
    even for the repair to exist. Doubling the whole certificate -- which
    changes nothing mathematically, since Q/den and G/den are unchanged --
    makes every residual even by construction.
    """
    Qint = [[int(round(Qf[i, j] * den)) for j in range(M)] for i in range(M)]
    # symmetrise exactly, then double
    for i in range(M):
        for j in range(i):
            v = (Qint[i][j] + Qint[j][i]) // 2
            Qint[i][j] = Qint[j][i] = v
    Qint = [[2 * v for v in row] for row in Qint]
    den, G = 2 * den, 2 * G

    # distinct representative (i, j) for each monomial alpha
    rep = {}
    for i, ei in enumerate(Z_BASIS):
        for j, ej in enumerate(Z_BASIS):
            a = tuple(ei[v] + ej[v] for v in range(N_VARS))
            if a not in rep:
                rep[a] = (i, j)

    acc = {}
    for i, ei in enumerate(Z_BASIS):
        for j, ej in enumerate(Z_BASIS):
            a = tuple(ei[v] + ej[v] for v in range(N_VARS))
            acc[a] = acc.get(a, 0) + Qint[i][j]

    zero = (0,) * N_VARS
    for a, (i, j) in rep.items():
        want = den * p_coeffs.get(a, 0) - (G if a == zero else 0)
        delta = want - acc[a]
        if delta == 0:
            continue
        if i == j:
            Qint[i][j] += delta            # contributes once
        else:
            assert delta % 2 == 0, "residual parity invariant violated"
            Qint[i][j] += delta // 2
            Qint[j][i] += delta // 2
    return Qint, den, G


def main():
    t0 = time.time()
    os.makedirs(OUT, exist_ok=True)
    log_lines = []

    def log(msg):
        line = f"[{time.time() - t0:7.1f}s] {msg}"
        print(line, flush=True)
        log_lines.append(line)

    p = load_p()
    log(f"p has {len(p)} terms; basis size {M}; monomials {len(P_MONOMIALS)}")

    A, alpha_index, _ = build_linear_map()
    S, AS, null, rank = sym_basis_nullspace(A)
    log(f"affine system rank {rank}, null-space dim {null.shape[1]}")

    rhs_p = np.zeros(len(P_MONOMIALS))
    for a, c in p.items():
        rhs_p[alpha_index[a]] = c
    rhs_1 = np.zeros(len(P_MONOMIALS))
    rhs_1[alpha_index[(0,) * N_VARS]] = 1.0

    Qp = solve_particular(AS, S, rhs_p)
    Q1 = solve_particular(AS, S, rhs_1)
    B = mats_from_sym(S, null)
    log(f"null-space blocks: {B.shape}")

    def feasible(gamma, t_start=None):
        Q0 = Qp - gamma * Q1
        t, lmin = max_lambda_min(Q0, B, t_start)
        return t, lmin

    log("bisection on gamma")
    lo, hi = -60.0, 60.0
    t_warm = None
    t_lo, l_lo = feasible(lo)
    log(f"  gamma={lo:+.4f}: lambda_min={l_lo:+.3e}")
    if l_lo <= 0:
        log("  lower bracket infeasible; widening")
        lo = -400.0
        t_lo, l_lo = feasible(lo)
        log(f"  gamma={lo:+.4f}: lambda_min={l_lo:+.3e}")
    t_warm = t_lo
    for it in range(40):
        mid = 0.5 * (lo + hi)
        t_mid, l_mid = feasible(mid, t_warm)
        if l_mid > 0:
            lo, t_warm = mid, t_mid
        else:
            hi = mid
        if it % 6 == 0 or hi - lo < 1e-7:
            log(f"  it {it}: gamma in [{lo:.8f}, {hi:.8f}], "
                f"lambda_min(mid)={l_mid:+.3e}")
        if hi - lo < 1e-9:
            break
    gamma_star = lo
    log(f"SOS optimum gamma* ~ {gamma_star:.10f}")

    # ---- retreat, re-centre, rationalise ----
    # NOTE: rationalise() doubles den (parity repair), so den_bits must be at
    # most 19 to respect the task's cap of 2^20.
    for back_off, den_bits in ((3e-6, 19), (5e-6, 19), (1e-5, 19),
                               (3e-5, 19), (1e-4, 19), (1e-3, 19)):
        gamma_t = gamma_star - back_off
        den = 1 << den_bits
        G = int(round(gamma_t * den))
        gamma_exact = G / den
        Q0 = Qp - gamma_exact * Q1
        t, lmin = max_lambda_min(Q0, B, t_warm)
        t, lmin2 = recenter(Q0, B, t)
        Qf = Q0 + np.tensordot(t, B, axes=(0, 0))
        lam = float(np.linalg.eigvalsh(Qf).min())
        log(f"back_off={back_off:g} den=2^{den_bits}: margin lambda_min={lam:.3e}")
        Qint, den_r, G_r = rationalise(Qf, den, p, G)
        ok = is_positive_definite(Qint)
        gt_num, gt_den = GAMMA_TARGET
        meets = G_r * gt_den >= gt_num * den_r
        log(f"  exact PD: {ok}; gamma={G_r / den_r:.10f}; meets target: {meets}")
        if ok and meets:
            den, G = den_r, G_r
            sub = {"den": den, "G": G, "Q": Qint}
            with open(os.path.join(OUT, "solution.json"), "w") as fh:
                json.dump(sub, fh)
            m = verify_submission(sub, p, GAMMA_TARGET, GAMMA_PARTIAL)
            score, reason = grade(m)
            log(f"self-check: score={score} ({reason}); gamma={G/den:.10f}")
            report = {"gamma_star_numeric": gamma_star, "gamma_certified": G / den,
                      "back_off": back_off, "den_bits": den_bits,
                      "margin_lambda_min": lam, "metrics": m,
                      "wall_seconds": time.time() - t0, "log": log_lines}
            with open(os.path.join(OUT, "solver_report.json"), "w") as fh:
                json.dump(report, fh, indent=2, default=float)
            log("done")
            return
    log("FAILED to produce an exactly-PD rational certificate")


if __name__ == "__main__":
    main()