"""Gauge-covariant Wilson-loop toy verification.

Run with: python wilson_experiment.py
Outputs results.json and prints quantitative predictions versus observations.
Quaternion q=(w,x,y,z) represents SU(2); Re Tr(U)/2 equals w.
"""
import json
import math
import numpy as np

SEED = 2472
rng = np.random.default_rng(SEED)


def qmul(a, b):
    a, b = np.asarray(a), np.asarray(b)
    w = a[..., 0]*b[..., 0] - np.sum(a[..., 1:]*b[..., 1:], axis=-1)
    v = (a[..., 0, None]*b[..., 1:] + b[..., 0, None]*a[..., 1:]
         + np.cross(a[..., 1:], b[..., 1:]))
    return np.concatenate([w[..., None], v], axis=-1)


def qconj(a):
    a = np.asarray(a).copy()
    a[..., 1:] *= -1
    return a


def qnorm(a):
    return a / np.linalg.norm(a, axis=-1, keepdims=True)


def random_q(n=1, angle=None):
    if angle is None:
        x = rng.normal(size=(n, 4))
        return qnorm(x)
    axis = rng.normal(size=(n, 3))
    axis /= np.linalg.norm(axis, axis=1, keepdims=True)
    a = np.broadcast_to(np.asarray(angle), (n,))
    return np.concatenate([np.cos(a)[:, None], np.sin(a)[:, None]*axis], axis=1)


def loop(uij, ujk, uki):
    return qmul(qmul(uij, ujk), uki)


def compatibility(qloops):
    return float(np.mean(qloops[:, 0]))


def gauge_transform(uij, hi, hj):
    return qmul(qmul(hi, uij), qconj(hj))


def analytic_mean(theta):
    # For independently isotropic fixed-angle edge errors, scalar loop mean is cos(theta)^3.
    return math.cos(theta)**3


def math_checks():
    n = 2000
    max_err = 0.0
    for _ in range(n):
        u, v, w = random_q(3)
        hi, hj, hk = random_q(3)
        h = loop(u, v, w)
        hp = loop(gauge_transform(u, hi, hj),
                  gauge_transform(v, hj, hk),
                  gauge_transform(w, hk, hi))
        expected = qmul(qmul(hi, h), qconj(hi))
        max_err = max(max_err, float(np.max(np.abs(hp - expected))))
    # A pure gauge connection U_ij=h_i h_j^dagger has exactly identity loops.
    flat = []
    for _ in range(1000):
        hi, hj, hk = random_q(3)
        flat.append(loop(qmul(hi, qconj(hj)), qmul(hj, qconj(hk)),
                         qmul(hk, qconj(hi))))
    flat = np.asarray(flat)
    return {"conjugation_max_abs_error": max_err,
            "pure_gauge_mean_M": compatibility(flat),
            "pure_gauge_max_energy": float(np.max(1-flat[:, 0]))}


def disorder_sweep():
    # Triangles are independent here; this isolates the predicted edge-disorder law.
    sigmas = np.linspace(0.0, 0.9, 10)
    rows = []
    for sigma in sigmas:
        vals = []
        for _ in range(300):
            angles = rng.normal(0.0, sigma, size=(300, 3))
            # Each edge has an independent random axis, but its scalar part is cos(angle).
            edges = [random_q(300, angles[:, i]) for i in range(3)]
            vals.append(np.mean(loop(edges[0], edges[1], edges[2])[:, 0]))
        vals = np.asarray(vals)
        observed = float(np.mean(vals))
        # E cos(theta) = exp(-sigma^2/2), so E M = exp(-3 sigma^2/2).
        predicted = math.exp(-1.5*sigma*sigma)
        rows.append({"sigma": float(sigma), "observed_M": observed,
                     "predicted_M": predicted,
                     "abs_error": abs(observed-predicted),
                     "susceptibility": float(300*np.var(vals))})
    return rows


def regularization_toy():
    """Optimize scalar edge angles to fit noisy target angles.

    Baseline fits each of three directed edge angles independently. Wilson adds
    lambda*(sum(edge angles))^2, the small-angle form of 1-ReTr(H)/2.
    This is a local coordinate chart of the SU(2) loop near identity.
    """
    targets = np.array([0.55, -0.40, 0.35])
    noise = np.array([0.30, -0.15, 0.20])
    y = targets + noise
    lr, steps = 0.08, 300
    out = []
    for lam in [0.0, 0.1, 0.5, 1.0]:
        x = np.zeros(3)
        for _ in range(steps):
            # task MSE + lambda*(1-cos(sum)); exact Wilson scalar in this chart
            s = np.sum(x)
            grad = 2*(x-y)/3.0 + lam*np.sin(s)
            x -= lr*grad
        task = float(np.mean((x-targets)**2))
        wilson = float(1-math.cos(np.sum(x)))
        out.append({"lambda": lam, "task_mse": task, "loop_energy": wilson,
                    "learned_angles": x.tolist()})
    return out


def main():
    checks = math_checks()
    sweep = disorder_sweep()
    opt = regularization_toy()
    # Quantitative acceptance tests: exact covariance, flatness, and disorder scaling.
    max_sweep_error = max(r["abs_error"] for r in sweep)
    result = {"seed": SEED, "math_checks": checks,
              "disorder_sweep": sweep, "regularization_toy": opt,
              "summary": {"max_disorder_prediction_error": max_sweep_error,
                          "disorder_prediction_tolerance": 0.015}}
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()