import json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

SEED = 2953
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device = "cuda" if torch.cuda.is_available() else "cpu"

# q(x|y) and r(y|x) are represented by two small categorical tables.
class TwoConditionals(nn.Module):
    def __init__(self, nx, ny):
        super().__init__()
        self.q_logits = nn.Parameter(torch.zeros(ny, nx))
        self.r_logits = nn.Parameter(torch.zeros(nx, ny))

    def log_q(self): return torch.log_softmax(self.q_logits, dim=1)
    def log_r(self): return torch.log_softmax(self.r_logits, dim=1)


def cycle_delta(logq, logr, x1, x2, y1, y2):
    # Indexing follows q[x|y] = logq[y,x], r[y|x] = logr[x,y].
    return (logq[y1, x1] + logr[x2, y1] + logq[y2, x2] + logr[x1, y2]
            - logr[x1, y1] - logq[y2, x1] - logr[x2, y2] - logq[y1, x2])


def all_quadruples(nx, ny, device):
    # Distinct pairs make the certificate nontrivial; all held-out evaluation
    # quadruples are deterministic rather than sampled from the training data.
    rows = []
    for x1 in range(nx):
      for x2 in range(nx):
       if x1 == x2: continue
       for y1 in range(ny):
        for y2 in range(ny):
         if y1 == y2: continue
         rows.append((x1,x2,y1,y2))
    return [torch.tensor([z[i] for z in rows], device=device) for i in range(4)]


def residuals(model, quads):
    lq, lr = model.log_q(), model.log_r()
    return cycle_delta(lq, lr, *quads).detach().cpu().numpy()


def train(joint, lam, steps=900):
    nx, ny = joint.shape
    model = TwoConditionals(nx, ny).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=0.08)
    p = torch.tensor(joint.ravel(), dtype=torch.float32, device=device)
    # Fixed synthetic population gives exactly matched data across conditions.
    rng = np.random.default_rng(SEED + 17)
    samples = rng.choice(nx*ny, size=5000, p=joint.ravel())
    xs = torch.tensor(samples // ny, dtype=torch.long, device=device)
    ys = torch.tensor(samples % ny, dtype=torch.long, device=device)
    train_x, train_y = xs[:4000], ys[:4000]
    qs = all_quadruples(nx, ny, device)
    # Hold out a deterministic quarter of quadruples for generalization testing.
    nq = qs[0].numel()
    perm = torch.randperm(nq, generator=torch.Generator(device=device).manual_seed(SEED+91), device=device)
    cut = int(0.75*nq)
    train_q = tuple(z[perm[:cut]] for z in qs)
    heldout_q = tuple(z[perm[cut:]] for z in qs)
    for step in range(steps):
        ix = torch.arange((step*64) % 3900, (step*64) % 3900 + 64, device=device) % 4000
        x, y = train_x[ix], train_y[ix]
        lq, lr = model.log_q(), model.log_r()
        task = -0.5 * (lq[y, x].mean() + lr[x, y].mean())
        # Use all small-support quadruples, which is cheap and lower variance.
        d = cycle_delta(lq, lr, *train_q)
        cycle = 0.5 * d.square().mean()
        loss = task + lam * cycle
        opt.zero_grad(); loss.backward(); opt.step()
    # Recover a joint from q and r using the least-squares log-ratio identity.
    with torch.no_grad():
        lq, lr = model.log_q(), model.log_r()
        # For a compatible pair, log p(x,y) differs from log q(x|y)+log p(y),
        # and p(y) can be estimated by averaging the two directional joints.
        joint_q = torch.exp(lq).T * torch.tensor(joint.sum(0), device=device)
        joint_r = torch.exp(lr) * torch.tensor(joint.sum(1), device=device)[:,None]
        recovered = 0.5*(joint_q + joint_r)
        recovered = recovered / recovered.sum()
        true = torch.tensor(joint, dtype=torch.float32, device=device)
        joint_l1 = torch.abs(recovered-true).sum().item()
        task_final = float((-0.5*(lq[train_y,train_x].mean()+lr[train_x,train_y].mean())).item())
    rr = residuals(model, qs)
    rh = residuals(model, heldout_q)
    return {"model": model, "p95_abs_residual": float(np.percentile(np.abs(rr),95)),
            "heldout_p95_abs_residual": float(np.percentile(np.abs(rh),95)),
            "mean_abs_residual": float(np.mean(np.abs(rr))),
            "heldout_mean_abs_residual": float(np.mean(np.abs(rh))), "joint_l1": joint_l1,
            "task_nll": task_final}


def main():
    # Strictly positive non-product joint, so the exact identity is meaningful.
    raw = np.array([[1.0, 2.0, 0.7, 1.4], [2.1, 0.8, 1.8, 0.5],
                    [0.6, 1.7, 2.4, 1.1], [1.3, 0.9, 1.5, 2.2]], dtype=np.float64)
    joint = raw/raw.sum()
    nx, ny = joint.shape
    px, py = joint.sum(1), joint.sum(0)
    q = joint/py[None,:]       # q[x,y]
    r = joint/px[:,None]       # r[x,y]
    # Core numerical verification: exact conditionals satisfy Delta=0.
    exact = []
    arbitrary = []
    rng = np.random.default_rng(SEED)
    aq = rng.dirichlet(np.ones(nx), size=ny)
    ar = rng.dirichlet(np.ones(ny), size=nx)
    for x1,x2,y1,y2 in [(0,1,0,1),(1,3,2,0),(2,3,1,3)]:
        exact.append(math.log(q[x1,y1])+math.log(r[x2,y1])+math.log(q[x2,y2])+math.log(r[x1,y2])
          -math.log(r[x1,y1])-math.log(q[x1,y2])-math.log(r[x2,y2])-math.log(q[x2,y1]))
        arbitrary.append(math.log(aq[y1,x1])+math.log(ar[x2,y1])+math.log(aq[y2,x2])+math.log(ar[x1,y2])
          -math.log(ar[x1,y1])-math.log(aq[y2,x1])-math.log(ar[x2,y2])-math.log(aq[y1,x2]))
    exact_max = float(np.max(np.abs(exact)))
    arbitrary_mean = float(np.mean(np.abs(arbitrary)))
    baseline = train(joint, 0.0)
    idea = train(joint, 0.5)
    out = {"seed":SEED,"device":device,"math_check_exact_max_abs_delta":exact_max,
           "math_check_random_mean_abs_delta":arbitrary_mean,"baseline":{k:v for k,v in baseline.items() if k!="model"},
           "idea_lambda_0.5":{k:v for k,v in idea.items() if k!="model"}}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    try:
        main()
    except (RuntimeError, torch.cuda.CudaError) as e:
        if device == "cuda":
            print("CUDA failed; rerun on CPU", repr(e))
            torch.cuda.empty_cache(); device = "cpu"; main()
        else: raise