import json, math, time
import torch
from quadratic_tree import make_random_tree, sequential_postorder, rake_levels, reconstruct, EdgeFactor

torch.set_num_threads(min(12, torch.get_num_threads()))
DT = torch.float64

def clone_case(case, requires=False):
    parent, U, h, k, edges = case
    def cp(x):
        y = x.detach().clone()
        y.requires_grad_(requires)
        return y
    ee = {c: EdgeFactor(cp(e.Hpp), cp(e.Hpc), cp(e.Hcc), cp(e.gp), cp(e.gc), cp(e.k))
          for c, e in edges.items()}
    return list(parent), cp(U), cp(h), cp(k), ee

def sync(device):
    if device.type == "cuda":
        torch.cuda.synchronize(device)

def math_checks():
    # Prediction 1: the eliminated parent quadratic correction is exactly quadratic
    # in coupling gamma: delta H / gamma^2 is constant.
    u, v = 1.7, 0.8
    gammas = [0.05, 0.1, 0.2, 0.4, 0.8]
    observed = [-(g*g)/(u+v) for g in gammas]
    predicted = [-(g*g)/(u+v) for g in gammas]
    rel_quad = max(abs(a-b) for a,b in zip(observed,predicted)) / max(abs(x) for x in predicted)
    ratios = [observed[i]/(gammas[i]*gammas[i]) for i in range(len(gammas))]
    ratio_spread = max(ratios)-min(ratios)

    # Prediction 2: damping changes the magnitude by (u+v)/(u+v+epsilon).
    gamma = 0.5
    epsilons = [0.0, 0.1, 0.3, 0.8]
    damp_obs = [abs(-(gamma*gamma)/(u+v+e)) for e in epsilons]
    damp_pred = [gamma*gamma/(u+v+e) for e in epsilons]
    damp_err = max(abs(a-b) for a,b in zip(damp_obs,damp_pred))
    monotone = all(damp_obs[i+1] < damp_obs[i] for i in range(len(damp_obs)-1))

    # Prediction 3: a balanced binary tree has logarithmic synchronized rake depth.
    depth_rows = []
    for n in [3, 7, 15, 31, 63]:
        parent = [-1] + [(i-1)//2 for i in range(1,n)]
        d = 2
        eye = torch.eye(d, dtype=DT)
        U = eye.unsqueeze(0).repeat(n,1,1)
        h = torch.zeros(n,d,dtype=DT); k = torch.zeros(n,dtype=DT)
        edges = {c: EdgeFactor(torch.zeros(d,d,dtype=DT), torch.zeros(d,d,dtype=DT),
                                eye, torch.zeros(d,dtype=DT), torch.zeros(d,dtype=DT),
                                torch.zeros((),dtype=DT)) for c in range(1,n)}
        _,_,_,_,levels,_ = rake_levels(parent,U,h,k,edges)
        depth_rows.append({"n": n, "observed_levels": len(levels),
                           "predicted_levels": int(math.floor(math.log2(n)))})
    depth_ok = all(x["observed_levels"] == x["predicted_levels"] for x in depth_rows)
    return {
        "quadratic_scaling": {"gammas": gammas, "observed_delta_H": observed,
                               "predicted_delta_H": predicted, "max_relative_error": rel_quad,
                               "delta_H_over_gamma2_spread": ratio_spread},
        "damping_scaling": {"epsilons": epsilons, "observed_magnitude": damp_obs,
                             "predicted_magnitude": damp_pred, "max_absolute_error": damp_err,
                             "strictly_decreasing": monotone},
        "balanced_depth": {"rows": depth_rows, "prediction_confirmed": depth_ok}
    }

def agreement(device):
    case = make_random_tree(31, 4, seed=77, device=device, dtype=DT)
    a = clone_case(case, requires=True)
    b = clone_case(case, requires=True)
    sr = sequential_postorder(*a, eps=1e-6)
    rr = rake_levels(*b, eps=1e-6)
    loss_s = sr[0].square().sum() + sr[1].square().sum() + sr[2].square()
    loss_r = rr[0].square().sum() + rr[1].square().sum() + rr[2].square()
    loss_s.backward(); loss_r.backward()
    # Inputs are cloned in the same order, so compare all differentiable tensors.
    max_out = max((sr[i]-rr[i]).abs().max().item() for i in range(3))
    max_grad = max((a[i].grad-b[i].grad).abs().max().item() for i in range(1,4))
    z = reconstruct(b[0], rr[5], b[1], b[2], b[3], b[4], rr[3], eps=1e-6)
    finite = bool(torch.isfinite(z).all())
    return {"max_root_coefficient_difference": max_out, "max_input_gradient_difference": max_grad,
            "reconstruction_finite": finite, "levels": len(rr[4])}

def benchmark(device):
    rows=[]
    for n in [127, 511, 1023]:
        case = make_random_tree(n, 8, seed=100+n, device=device, dtype=DT)
        times=[]
        for mode in ["sequential", "rake"]:
            # Warmup and three measured runs; no autograd in timing.
            fn = sequential_postorder if mode == "sequential" else rake_levels
            fn(*case, eps=1e-6)
            sync(device); t0=time.perf_counter()
            result=None
            for _ in range(3): result=fn(*case, eps=1e-6)
            sync(device); elapsed=(time.perf_counter()-t0)/3
            times.append(elapsed*1000)
        rows.append({"n":n, "sequential_ms":times[0], "rake_ms":times[1],
                     "rake_over_sequential":times[1]/times[0]})
    return rows

def main():
    try:
        device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # CUDA errors, including shared-memory exhaustion, must not abort the run.
        math_result=math_checks()
        agree=agreement(device)
        bench=benchmark(device)
    except Exception as exc:
        device=torch.device("cpu")
        math_result=math_checks()
        agree=agreement(device)
        bench=benchmark(device)
        fallback=str(exc)
    else:
        fallback=None
    result={"device":str(device), "math_checks":math_result, "agreement":agree,
            "benchmark":bench, "cuda_fallback_error":fallback}
    with open("results.json","w") as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__ == "__main__": main()