import json, math, time
from pathlib import Path
import numpy as np
import torch

# Cross-Partial Nash Compatibility Regularizer MVP.
# q1=a*u1*u2 and q2=b*u1*u2 are deliberately incompatible unless a=b.
# The supervised target (a,b)=(1,-1) creates a controlled critic-fitting problem.

torch.set_default_dtype(torch.float64)
SEED = 1167

def set_seed(seed=SEED):
    np.random.seed(seed); torch.manual_seed(seed)


def cross_residual(a, b, u):
    """Autograd implementation of d/du2(dq1/du1)-d/du1(dq2/du2)."""
    u = u.detach().clone().requires_grad_(True)
    q1 = a * u[:, 0] * u[:, 1]
    q2 = b * u[:, 0] * u[:, 1]
    g1 = torch.autograd.grad(q1.sum(), u, create_graph=True)[0][:, 0]
    g2 = torch.autograd.grad(q2.sum(), u, create_graph=True)[0][:, 1]
    h12 = torch.autograd.grad(g1.sum(), u, create_graph=True)[0][:, 1]
    h21 = torch.autograd.grad(g2.sum(), u, create_graph=True)[0][:, 0]
    return h12 - h21


def analytic_lambda_prediction(lam):
    # argmin .5[(a-1)^2+(b+1)^2] + lam*(a-b)^2 gives d=a-b=2/(1+4 lam).
    return 2.0 / (1.0 + 4.0 * lam)


def solve_coeff(lam, steps=500, lr=0.04):
    set_seed()
    ab = torch.tensor([0.0, 0.0], requires_grad=True)
    opt = torch.optim.SGD([ab], lr=lr)
    u = torch.linspace(-1, 1, 64).reshape(-1, 1).repeat(1, 2)
    target = torch.tensor([1.0, -1.0])
    for _ in range(steps):
        # The coefficient objective is the exact average critic fitting objective.
        fit = 0.5 * ((ab - target) ** 2).sum()
        compat = (ab[0] - ab[1]) ** 2
        loss = fit + lam * compat
        opt.zero_grad(); loss.backward(); opt.step()
    residual = float((ab[0] - ab[1]).detach().abs())
    fit = float((0.5 * ((ab.detach() - target) ** 2).sum()))
    return ab.detach().numpy().tolist(), residual, fit


def gamma_sweep():
    # Gibbs log-conditionals have cross derivative mismatch (a-b)/gamma.
    # Therefore the squared mismatch must be residual^2/gamma^2.
    d = 2.0
    out = []
    for gamma in [0.25, 0.5, 1.0, 2.0, 4.0]:
        observed = (d / gamma) ** 2
        predicted = 4.0 / gamma**2
        out.append({"gamma": gamma, "observed": observed, "predicted": predicted,
                    "relative_error": abs(observed-predicted)/(predicted+1e-12)})
    return out


def lambda_sweep():
    out = []
    for lam in [0.0, 0.05, 0.25, 0.5, 1.0, 2.0, 5.0]:
        ab, observed, fit = solve_coeff(lam)
        predicted = analytic_lambda_prediction(lam)
        out.append({"lambda": lam, "a_b": ab, "observed_abs_residual": observed,
                    "predicted_abs_residual": predicted, "fit_loss": fit,
                    "relative_error": abs(observed-predicted)/(predicted+1e-12)})
    return out


def finite_difference_check():
    # Independently check the autograd residual against finite differences.
    u = torch.tensor([[0.37, -0.21]])
    a, b = 1.3, -0.7
    auto = float(cross_residual(torch.tensor(a), torch.tensor(b), u)[0])
    eps = 1e-5
    def q1(x, y): return a*x*y
    def q2(x, y): return b*x*y
    # d/du2 dq1/du1 and d/du1 dq2/du2 via central differences of first derivatives.
    fd12 = (q1(1, u[0,1].item()+eps)-q1(1, u[0,1].item()-eps))/(2*eps)
    fd21 = (q2(u[0,0].item()+eps, 1)-q2(u[0,0].item()-eps, 1))/(2*eps)
    return {"autograd_residual": auto, "finite_difference_residual": fd12-fd21,
            "absolute_error": abs(auto-(fd12-fd21))}


def mini_compare():
    # Same coefficient critic and optimizer, interpreted as a tiny centralized critic.
    rows = []
    for lam in [0.0, 1.0]:
        ab, residual, fit = solve_coeff(lam, steps=500, lr=0.04)
        rows.append({"method": "baseline" if lam == 0 else "compat_regularized",
                     "lambda": lam, "final_abs_cross_partial": residual,
                     "critic_fit_loss": fit, "coefficients": ab})
    return rows


def main():
    set_seed()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    # This MVP is deliberately CPU-sized; no GPU allocation is needed.
    result = {"device_available": device, "finite_difference": finite_difference_check(),
              "lambda_sweep": lambda_sweep(), "gamma_sweep": gamma_sweep(),
              "mini_compare": mini_compare(),
              "predictions": [
                "At lambda=0 residual=2; increasing lambda follows 2/(1+4 lambda).",
                "The Gibbs incompatibility squared scales as 1/gamma^2, so halving gamma quadruples it.",
                "At large lambda residual tends to zero (compatibility), while fitting error increases because targets are intentionally incompatible."
              ]}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()