import json
import time
import numpy as np


def energy(S, mu, v):
    return float(-0.5 * v @ S @ v - mu @ v)


def reduce_blocks(S, mu, v_H, H, Q):
    """Exact conditional objective represented by the active block and field."""
    H = np.asarray(H, dtype=int); Q = np.asarray(Q, dtype=int)
    return S[np.ix_(Q, Q)], mu[Q] + S[np.ix_(Q, H)] @ v_H


def certified_candidates(S, mu, v, H, Q, eps=1e-12):
    """Return active indices certified for every assignment of the other active spins."""
    H = np.asarray(H, dtype=int); Q = np.asarray(Q, dtype=int)
    out = []
    for i in Q:
        fixed = mu[i] + (S[i, H] @ v[H] if len(H) else 0.0)
        unresolved_bound = np.abs(S[i, Q]).sum() - abs(S[i, i])
        if abs(fixed) > unresolved_bound + eps and np.sign(v[i]) == np.sign(fixed):
            out.append(int(i))
    return out


def verify_math(seed=7):
    rng = np.random.default_rng(seed)
    n = 7
    A = rng.normal(size=(n, n)); S = (A + A.T) / 2; np.fill_diagonal(S, 0)
    mu = rng.normal(size=n); v = rng.choice([-1, 1], size=n)
    # Exact conditional identity for several arbitrary H/Q assignments.
    H = np.array([0, 3, 5]); Q = np.array([1, 2, 4, 6])
    block, meff = reduce_blocks(S, mu, v[H], H, Q)
    C = energy(S, mu, v) - (-0.5 * v[Q] @ block @ v[Q] - meff @ v[Q])
    identity_max_err = 0.0
    for bits in range(1 << len(Q)):
        q = np.array([1 if bits >> k & 1 else -1 for k in range(len(Q))])
        full = np.zeros(n, dtype=int); full[H] = v[H]; full[Q] = q
        reduced = C - 0.5 * q @ block @ q - meff @ q
        identity_max_err = max(identity_max_err, abs(energy(S, mu, full) - reduced))
    # Construct a guaranteed certified case and exhaustively check all Q assignments.
    m = 6; B = rng.normal(size=(m, m)); T = (B + B.T) / 2; np.fill_diagonal(T, 0)
    b = np.zeros(m); b[0] = 5.0; b[1] = -4.5
    state = np.ones(m, dtype=int); state[1] = -1
    H2 = np.array([], dtype=int); Q2 = np.arange(m)
    cand = certified_candidates(T, b, state, H2, Q2)
    violations = 0
    for i in cand:
        for bits in range(1 << len(Q2)):
            q = np.array([1 if bits >> k & 1 else -1 for k in range(len(Q2))])
            if state[i] != np.sign(b[i] + T[i] @ q):
                violations += 1
    return {"conditional_identity_max_abs_error": identity_max_err,
            "certified_indices": cand, "certification_violations": violations}


def run_solver(S, mu, v0, steps, reduce=False, eps=1e-10):
    n = len(mu); v = v0.copy(); frozen = np.zeros(n, dtype=bool)
    active = np.arange(n); active_bias = mu.copy(); work = 0
    cert_viol = 0; active_trace = []
    t0 = time.perf_counter()
    for _ in range(steps):
        if reduce:
            # Certification uses the exact current frozen field and remaining tail.
            H = np.flatnonzero(frozen); Q = np.flatnonzero(~frozen)
            cand = certified_candidates(S, mu, v, H, Q, eps)
            for i in cand:
                frozen[i] = True
            active = np.flatnonzero(~frozen)
            if len(active):
                # This is the same induced field as incremental maintenance.
                active_bias = mu[active] + S[np.ix_(active, np.flatnonzero(frozen))] @ v[frozen]
        if len(active):
            fields = active_bias + S[np.ix_(active, active)] @ v[active]
            v[active] = np.where(fields >= 0, 1, -1)
            work += len(active) * len(active)
        active_trace.append(int(len(active)))
        # Count post-update violations for frozen spins against the actual full state.
        if reduce and frozen.any():
            h = mu + S @ v
            cert_viol += int(np.sum(frozen & (v * h <= 0)))
    return {"energy": energy(S, mu, v), "state": v, "seconds": time.perf_counter()-t0,
            "work_matvec_entries": int(work), "active_trace": active_trace,
            "frozen": int(frozen.sum()), "certification_violations": cert_viol}


def benchmark(seed=11, n=1200, steps=35):
    rng = np.random.default_rng(seed)
    A = rng.normal(size=(n, n)).astype(np.float64)
    S = (A + A.T) * (0.5 / np.sqrt(n)); np.fill_diagonal(S, 0)
    # Strong polarized coordinates make the claimed late-stage tail reduction visible,
    # while the remaining coordinates retain a nontrivial dense Ising problem.
    mu = rng.normal(0, 0.10, n); strong = int(0.70*n)
    mu[:strong] += 30.0 * rng.choice([-1, 1], strong)
    v0 = rng.choice([-1, 1], n)
    base = run_solver(S, mu, v0, steps, reduce=False)
    idea = run_solver(S, mu, v0, steps, reduce=True)
    return {"n": n, "steps": steps, "strong_field_fraction": strong/n,
            "baseline": {k:v for k,v in base.items() if k != "state"},
            "idea": {k:v for k,v in idea.items() if k != "state"},
            "final_energy_difference_idea_minus_baseline": idea["energy"]-base["energy"],
            "work_ratio_idea_over_baseline": idea["work_matvec_entries"]/base["work_matvec_entries"],
            "final_state_agreement": float(np.mean(idea["state"] == base["state"]))}


if __name__ == "__main__":
    result = {"math": verify_math(), "benchmark": benchmark()}
    print(json.dumps(result, indent=2))