import json, math, random
import numpy as np
import torch

SEED = 1737
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_default_dtype(torch.float64)


def exact_and_hutchinson(theta, A, s, direction, beta=0.0, rmin=1e-3, probes=256):
    """Compute D_w for rho=rmin+softplus(s+direction dot theta), and Hutchinson div(G)."""
    theta = theta.detach().requires_grad_(True)
    s = s.detach().requires_grad_(True)
    L = 0.5 * theta @ A @ theta
    g = torch.autograd.grad(L, theta, create_graph=True)[0]
    rho = rmin + torch.nn.functional.softplus(s + direction @ theta)
    G = -rho * g
    div = 0.0
    for i in range(theta.numel()):
        div = div + torch.autograd.grad(G[i], theta, retain_graph=True, create_graph=True)[0][i]
    D = div + beta * rho * (g @ g)
    vals = []
    for _ in range(probes):
        z = torch.randint(0, 2, theta.shape, dtype=theta.dtype) * 2 - 1
        jgz = torch.autograd.grad((G * z).sum(), theta, retain_graph=True)[0]
        vals.append((z * jgz).sum().item())
    return float(D.detach()), float(div.detach()), float(rho.detach()), np.array(vals)


def verify():
    A = torch.diag(torch.tensor([1.0, 7.0, 30.0]))
    theta = torch.tensor([0.7, -0.4, 0.2])
    s = torch.tensor(-0.3)
    direction = torch.tensor([0.13, -0.21, 0.17])
    D, div, rho, h = exact_and_hutchinson(theta, A, s, direction, beta=.08, probes=512)
    # Explicit formula: div(-rho*g)=-rho tr(A)-grad(rho).g.
    with torch.no_grad():
        u = s + direction @ theta
        drho_dot_g = torch.sigmoid(u) * (direction @ (A @ theta))
        explicit_div = -rho * torch.trace(A).item() - drho_dot_g.item()
        explicit_D = explicit_div + .08 * rho * (theta @ A @ A @ theta).item()
    return {
        "exact_D": D, "formula_D": explicit_D, "abs_D_error": abs(D-explicit_D),
        "exact_div": div, "formula_div": explicit_div,
        "hutch_mean": float(h.mean()), "hutch_std_error": float(h.std(ddof=1)/math.sqrt(len(h))),
        "rho": rho
    }


def run_optimizer(adaptive, seed, steps=180):
    torch.manual_seed(seed); np.random.seed(seed)
    A = torch.diag(torch.tensor([1.0, 15.0]))
    theta = torch.tensor([2.0, 2.0], dtype=torch.float64)
    # A fixed directional dependence makes grad rho nonzero, while s is the controller.
    direction = torch.tensor([0.25, -0.20])
    s = torch.tensor(-2.0, requires_grad=True)
    s_lr = .08
    fixed_rho = 0.055
    losses=[]; rhos=[]; contractions=[]; updates=[]
    for t in range(steps):
        theta.requires_grad_(True)
        L = .5 * theta @ A @ theta
        g = torch.autograd.grad(L, theta, create_graph=True)[0]
        if adaptive:
            rho = .002 + torch.nn.functional.softplus(s + direction @ theta)
            G = -rho * g
            div = sum(torch.autograd.grad(G[i], theta, retain_graph=True, create_graph=True)[0][i] for i in range(2))
            D = div  # beta=0: weighted divergence is ordinary divergence
            # Controller minimizes a soft contraction-margin violation, but does not alter theta gradient.
            penalty = torch.nn.functional.softplus(D + 0.35)
            ds = torch.autograd.grad(penalty, s, allow_unused=True)[0]
            if ds is not None:
                with torch.no_grad(): s -= s_lr * ds.clamp(-2,2)
            step = (-rho.detach() * g.detach())
            rho_value=float(rho.detach())
            D_value=float(D.detach())
        else:
            step = -fixed_rho * g.detach()
            rho_value=fixed_rho
            D_value=-fixed_rho * torch.trace(A).item()
        theta = (theta.detach() + step)
        losses.append(float(L.detach())); rhos.append(rho_value); contractions.append(D_value); updates.append(float(step.norm()))
    return {"final_loss": losses[-1], "initial_loss": losses[0], "losses": losses,
            "mean_rho":float(np.mean(rhos)), "rho_last":rhos[-1],
            "mean_D":float(np.mean(contractions)), "max_loss":float(max(losses)),
            "mean_update":float(np.mean(updates)), "final_theta":theta.tolist()}


def main():
    check=verify()
    base=[run_optimizer(False, 100+i) for i in range(8)]
    idea=[run_optimizer(True, 100+i) for i in range(8)]
    def summary(xs):
        return {"final_loss_mean":float(np.mean([x['final_loss'] for x in xs])),
                "final_loss_std":float(np.std([x['final_loss'] for x in xs], ddof=1)),
                "max_loss_mean":float(np.mean([x['max_loss'] for x in xs])),
                "mean_D":float(np.mean([x['mean_D'] for x in xs])),
                "mean_rho":float(np.mean([x['mean_rho'] for x in xs]))}
    out={"verification":check, "baseline":summary(base), "idea":summary(idea),
         "settings":{"steps":180,"seeds":8,"A":[1,15],"beta":0,"baseline_rho":.055}}
    with open("results.json","w") as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__ == '__main__': main()