"""Gauge-Patched Local Experts: toy math checks and a tiny optimization study.

Run: python3 gauge_patched_experts.py
Outputs: results.json and a concise terminal report.
"""
import json, math
from pathlib import Path
import numpy as np
import torch

SEED = 1270
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.set_num_threads(4)


def rot(theta):
    c, s = np.cos(theta), np.sin(theta)
    return np.array([[c, -s], [s, c]], dtype=float)


def holder(vals, xs, beta):
    # Dense pairwise Holder seminorm, adequate for this 1D sanity check.
    dx = np.abs(xs[:, None] - xs[None, :])
    dv = np.abs(vals[:, None] - vals[None, :])
    mask = dx > 1e-12
    return float(np.max(dv[mask] / (dx[mask] ** beta)))


def math_checks():
    beta = 0.75
    # Prediction 1: exp(skew) is orthogonal, hence preserves every feature norm.
    norm_errors = []
    for theta in np.linspace(-3.0, 3.0, 31):
        h = np.array([0.7, -1.3])
        g = rot(theta)  # exact SO(2) exponential
        norm_errors.append(abs(np.linalg.norm(g @ h) - np.linalg.norm(h)))
    p1 = {"prediction": "orthogonal transition preserves feature norm", "max_abs_norm_error": max(norm_errors)}

    # Prediction 2: a transition extended from identity across a gap d has
    # Holder size proportional to d^{-beta}. Use a smooth cubic bump.
    theta = 0.4
    ds = np.array([0.04, 0.06, 0.09, 0.13, 0.19, 0.28])
    hs = []
    xs = np.linspace(-1.0, 1.0, 1601)
    for d in ds:
        # transition changes only across [0,d], identity before and rot(theta) after
        t = np.clip(xs / d, 0, 1)
        bump = t*t*(3 - 2*t)
        # angle field; Holder norm of matrix entries is equivalent up to constants
        vals = np.array([rot(theta*a)[0, 0] for a in bump])
        hs.append(holder(vals, xs, beta))
    slope = float(np.polyfit(np.log(ds), np.log(hs), 1)[0])
    p2 = {"prediction": "extension Holder norm scales as gap^{-beta}", "beta": beta,
          "observed_loglog_slope": slope, "predicted_slope": -beta,
          "gaps": ds.tolist(), "holder_norms": [float(x) for x in hs]}

    # Prediction 3: for fixed h, ||h - R(eps)h||^2 is quadratic in eps.
    h = np.array([1.2, -0.8])
    eps = np.array([0.005, 0.01, 0.02, 0.04, 0.08, 0.16])
    residuals = np.array([np.linalg.norm(h - rot(e) @ h)**2 for e in eps])
    slope3 = float(np.polyfit(np.log(eps), np.log(residuals), 1)[0])
    ratios = residuals / (eps**2 * np.dot(h, h))
    p3 = {"prediction": "small-angle patch residual is quadratic", "observed_loglog_slope": slope3,
          "predicted_slope": 2.0, "residual_over_eps2_norm2": [float(x) for x in ratios]}
    return [p1, p2, p3]


def tiny_experiment():
    """Two local experts see differently rotated coordinates of the same target.
    Both variants have the same trainable expert offsets and task loss. The idea
    adds the overlap alignment penalty using the known orthogonal transition.
    """
    torch.manual_seed(SEED)
    n = 512
    z = torch.randn(n, 2)
    true_w = torch.tensor([[1.0, -0.4], [0.3, 0.8]])
    y = z @ true_w.T
    angle = 0.9
    G = torch.tensor(rot(angle), dtype=torch.float32)
    # Expert 0 receives z, expert 1 receives G z; target is shared latent output.
    x0, x1 = z, z @ G.T
    torch.manual_seed(SEED + 77)
    init0 = torch.randn(2,2)*0.5
    init1 = torch.randn(2,2)*0.5
    def run(lam, steps=500):
        # Matched initialization isolates the effect of patch regularization.
        W0 = torch.nn.Parameter(init0.clone())
        W1 = torch.nn.Parameter(init1.clone())
        opt = torch.optim.Adam([W0,W1], lr=0.04)
        for _ in range(steps):
            h0, h1 = x0 @ W0.T, x1 @ W1.T
            task = ((h0-y)**2).mean() + ((h1-y)**2).mean()
            # h0 ~= G^{-1} h1, equivalently h0 ~= G.T h1.
            patch = ((h0 - h1 @ G)**2).mean()
            loss = task + lam * patch
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            h0, h1 = x0 @ W0.T, x1 @ W1.T
            task = float((((h0-y)**2).mean()+((h1-y)**2).mean()).item()/2)
            disagreement = float(((h0-h1@G)**2).mean().item())
        return task, disagreement
    lambdas = [0.0, 0.01, 0.1, 1.0, 10.0]
    sweep = []
    for lam in lambdas:
        task, disagreement = run(lam)
        sweep.append({"lambda_patch": lam, "task_mse": task, "overlap_mse": disagreement})
    return {"metric": "mean task MSE; aligned overlap MSE", "matched_initialization": True,
            "baseline": sweep[0], "gauge_patched_lambda_1": sweep[3], "lambda_sweep": sweep}


def main():
    checks = math_checks()
    exp = tiny_experiment()
    out = {"seed": SEED, "math_checks": checks, "tiny_experiment": exp}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()