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

SEED = 19
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
try:
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    if device == 'cuda': torch.zeros(1, device='cuda')
except Exception:
    device = 'cpu'


def cube_neighbors(n, dev):
    ids = torch.arange(2 ** n, device=dev)
    return torch.stack([ids ^ (1 << i) for i in range(n)], 1)


def laplacian_target(f, nb):
    return 0.5 * (f[:, None] - f[nb]).sum(1)


def divergence(v, nb):
    n = nb.shape[1]
    return (v - v[nb, torch.arange(n, device=v.device)]).sum(1)


def solve_beckmann(f, nb, steps=80, rho=2.0, eta=0.002, eps=1e-5):
    """Differentiable unrolled minimization of mean(||V||_2)+rho*mean(div(V)-Lf)^2."""
    n = nb.shape[1]; N = f.numel()
    b = laplacian_target(f, nb)
    v = torch.zeros((N, n), device=f.device, dtype=f.dtype)
    cols = torch.arange(n, device=f.device)
    for _ in range(steps):
        r = divergence(v, nb) - b
        # div^* r has entries r(x)-r(x^flip); div is self-adjoint here.
        adj = r[:, None] - r[nb]
        norm = torch.sqrt((v * v).sum(1, keepdim=True) + eps)
        # adj has shape [N,n]; this is the gradient of the residual term.
        grad = v / norm + 2.0 * rho * adj
        v = v - eta * grad
    residual = divergence(v, nb) - b
    value = torch.sqrt((v * v).sum(1) + eps).mean() + rho * (residual ** 2).mean()
    return value, v, residual, b


class SmallMLP(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(n, 24), nn.Tanh(), nn.Linear(24, 1))
    def forward(self, x): return self.net(x).squeeze(1)


def metrics(model, x, y, nb, train_idx):
    with torch.no_grad():
        f = model(x); pred = (f > 0).long()
        acc = (pred == y).float().mean().item()
        one = (f[:, None] - f[nb]).abs().mean().item()
        # random reproducible two-bit perturbation, using coordinates 0 and 1
        two_nb = nb[:, 0][nb[:, 1]] if False else None
        ids = torch.arange(x.shape[0], device=x.device)
        two = ids ^ 1 ^ 2
        multi = (f - f[two]).abs().mean().item()
        ce = nn.functional.binary_cross_entropy_with_logits(f[train_idx], y[train_idx].float()).item()
    return {'accuracy_all': acc, 'one_bit_logit_change': one,
            'two_bit_logit_change': multi, 'train_ce': ce}


def run_experiment(n=8, train_fraction=0.5, steps=260):
    nb = cube_neighbors(n, device)
    N = 2 ** n
    ids = torch.arange(N, device=device)
    # {-1,+1} cube; bit order is consistent with neighbor table.
    x = 2.0 * ((ids[:, None] >> torch.arange(n, device=device)) & 1).float() - 1.0
    # parity interaction on first three coordinates, with a weak nuisance term.
    y = (((x[:, :3].prod(1) > 0)).long())
    perm = torch.randperm(N, device=device)
    train_idx = perm[:int(train_fraction * N)]
    out = {}
    for name in ('ce', 'edge', 'beckmann'):
        torch.manual_seed(SEED + {'ce': 0, 'edge': 1, 'beckmann': 2}[name])
        model = SmallMLP(n).to(device)
        opt = torch.optim.Adam(model.parameters(), lr=0.025)
        last_flow = None
        for t in range(steps):
            opt.zero_grad()
            f = model(x)
            loss = nn.functional.binary_cross_entropy_with_logits(f[train_idx], y[train_idx].float())
            if name == 'edge':
                edge = ((f[:, None] - f[nb]) ** 2).mean()
                loss = loss + 0.018 * edge
            elif name == 'beckmann':
                br, _, res, _ = solve_beckmann(f, nb)
                loss = loss + 0.008 * br
                last_flow = res.detach()
            loss.backward(); opt.step()
        m = metrics(model, x, y, nb, train_idx)
        if name == 'beckmann':
            with torch.no_grad():
                _, _, rr, bb = solve_beckmann(model(x), nb)
            m['flow_residual_rms'] = float(torch.sqrt((rr * rr).mean()).cpu())
            m['target_rms'] = float(torch.sqrt((bb * bb).mean()).cpu())
        m['steps'] = steps
        out[name] = m
    return out


def verify_math(n=4):
    nb = cube_neighbors(n, device); N = 2 ** n
    g = torch.randn(N, device=device)
    v = torch.randn(N, n, device=device)
    b = laplacian_target(g, nb)
    d = divergence(v, nb)
    # Both operators have zero mean, and divergence is the adjoint of itself.
    mean_b = abs(float(b.mean()))
    mean_d = abs(float(d.mean()))
    lhs = (divergence(v, nb) * g).sum()
    rhs = (v * (g[:, None] - g[nb])).sum()
    adjoint_err = abs(float(lhs - rhs))
    # A convex solver should reduce its penalized objective from V=0.
    f = torch.sin(torch.arange(N, device=device).float())
    b_f = laplacian_target(f, nb)
    # Explicit feasible reference: V_i=D_i f/2 gives div(V)=Lf exactly.
    vf = 0.25 * (f[:, None] - f[nb])
    feasible_residual = divergence(vf, nb) - b_f
    zero_obj = 2.0 * (b_f ** 2).mean()
    val, vv, rr, _ = solve_beckmann(f, nb)
    return {'mean_abs_Lf': mean_b, 'mean_abs_divV': mean_d,
            'adjoint_relative_error': adjoint_err / (abs(float(lhs)) + 1e-8),
            'explicit_feasible_flow_residual': float(torch.sqrt((feasible_residual ** 2).mean())),
            'explicit_feasible_flow_norm': float(torch.sqrt((vf * vf).sum(1)).mean()),
            'zero_flow_penalized_objective': float(zero_obj),
            'optimized_penalized_objective': float(val.detach()),
            'solver_residual_rms': float(torch.sqrt((rr * rr).mean()).detach()),
            'solver_residual_over_target_rms': float(torch.sqrt((rr * rr).mean()) / (torch.sqrt((b_f * b_f).mean()) + 1e-8))}


if __name__ == '__main__':
    checks = verify_math()
    results = run_experiment()
    report = {'device': device, 'math_check': checks, 'experiment': results}
    with open('results.json', 'w') as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))