"""
Empirical Bayes analytical detection head.

Bayesian linear regression: W ~ N(0, tau^2 I), Y|X,W ~ N(XW, sigma^2 I)
The optimal regularization is lambda = sigma^2 / tau^2.

Empirical Bayes estimates sigma^2 and tau^2 from the data by maximizing
the log marginal likelihood (type-II ML):

  log p(Y|X, sigma^2, tau^2) = -n/2 log(2pi) - 1/2 log|sigma^2 I + tau^2 X X^T|
                                 - 1/2 Y^T (sigma^2 I + tau^2 X X^T)^{-1} Y

Using the SVD of X = U S V^T, this simplifies to operations on the singular values.
The optimization alternates between updating sigma^2 and tau^2.

This gives a principled, per-task lambda with calibrated uncertainty.
"""

import json, os, sys, time
import torch
import torch.nn.functional as F

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
CACHE_DIR = os.environ.get("ARENA_CACHE_DIR")
COCO_ROOT = os.environ.get("ARENA_COCO_ROOT")
VAL_CACHE = os.environ.get("ARENA_VAL_CACHE")
DEVICE = "cuda"
RESOLUTION = 640
NUM_CLASSES = 80


def cofiber_decompose(f, n_scales):
    cofibers = []; residual = f
    for _ in range(n_scales - 1):
        omega = F.avg_pool2d(residual, 2)
        sigma_omega = F.interpolate(omega, size=residual.shape[2:], mode="bilinear", align_corners=False)
        cofibers.append(residual - sigma_omega); residual = omega
    cofibers.append(residual); return cofibers


def make_locations(sizes, strides, device="cpu"):
    locs = []
    for (h, w), s in zip(sizes, strides):
        ys = (torch.arange(h, device=device, dtype=torch.float32) + 0.5) * s
        xs = (torch.arange(w, device=device, dtype=torch.float32) + 0.5) * s
        gy, gx = torch.meshgrid(ys, xs, indexing="ij")
        locs.append(torch.stack([gx.flatten(), gy.flatten()], -1))
    return locs


def assign_targets(loc, boxes, labels, stride, sr):
    n = loc.shape[0]
    ct = torch.full((n,), -1, dtype=torch.long); rt = torch.zeros(n, 4); ctrt = torch.zeros(n)
    if boxes.numel() == 0: return ct, rt, ctrt
    areas = (boxes[:,2]-boxes[:,0])*(boxes[:,3]-boxes[:,1])
    l=loc[:,None,0]-boxes[None,:,0]; t=loc[:,None,1]-boxes[None,:,1]
    r=boxes[None,:,2]-loc[:,None,0]; b=boxes[None,:,3]-loc[:,None,1]
    ltrb=torch.stack([l,t,r,b],-1); in_box=ltrb.min(-1).values>0
    cx=(boxes[:,0]+boxes[:,2])/2; cy=(boxes[:,1]+boxes[:,3])/2; rad=stride*1.5
    in_center=((loc[:,None,0]>=cx-rad)&(loc[:,None,0]<=cx+rad)&(loc[:,None,1]>=cy-rad)&(loc[:,None,1]<=cy+rad))
    max_d=ltrb.max(-1).values; in_level=(max_d>=sr[0])&(max_d<=sr[1])
    pos=in_box&in_center&in_level; a=areas[None,:].expand_as(pos).clone(); a[~pos]=float("inf")
    matched=a.argmin(1); is_pos=a.gather(1,matched[:,None]).squeeze(1)<float("inf")
    ct[is_pos]=labels[matched[is_pos]]
    if is_pos.any():
        rt[is_pos]=ltrb[torch.arange(n)[is_pos],matched[is_pos]]
        lp,tp,rp,bp=rt[is_pos].unbind(-1)
        ctrt[is_pos]=torch.sqrt((torch.minimum(lp,rp)/torch.maximum(lp,rp).clamp(min=1e-6))*(torch.minimum(tp,bp)/torch.maximum(tp,bp).clamp(min=1e-6)))
    return ct, rt, ctrt


def empirical_bayes_lambda(XtX, XtY, n_samples, n_iters=50):
    """Estimate optimal lambda = sigma^2/tau^2 via empirical Bayes.

    Uses the eigendecomposition of XtX to iteratively estimate:
      tau^2 = (1/d) sum_i (gamma_i * w_i^2)  where gamma_i = s_i^2 / (s_i^2 + lambda)
      sigma^2 = (1/(n - sum gamma_i)) * ||Y - X W||^2

    gamma_i is the "effective number of well-determined parameters" (Mackay).
    """
    d = XtX.shape[0]
    k = XtY.shape[1]  # number of output dims

    # Eigendecompose XtX
    eigvals, eigvecs = torch.linalg.eigh(XtX)
    eigvals = eigvals.clamp(min=1e-10)

    # Project targets into eigenspace
    VtXtY = eigvecs.T @ XtY  # (d, k)

    # Initialize
    alpha = 1.0   # 1/tau^2
    beta = 1.0    # 1/sigma^2
    lam = alpha / beta

    for it in range(n_iters):
        # Effective parameters per eigenvalue
        gamma = eigvals / (eigvals + lam)
        gamma_sum = gamma.sum().item()

        # Solve with current lambda
        W = eigvecs @ (VtXtY / (eigvals + lam).unsqueeze(1))

        # Residual sum of squares (averaged over output dims)
        # ||Y - XW||^2 = ||Y||^2 - 2 W^T X^T Y + W^T X^T X W
        # = YtY - 2 W^T XtY + W^T XtX W
        # But we don't have YtY. Approximate from XtY and XtX.
        # W^T XtY = sum of VtXtY_i^2 * s_i / (s_i + lam)
        WtXtY = (VtXtY ** 2 * eigvals.unsqueeze(1) / (eigvals + lam).unsqueeze(1) ** 2).sum(0)

        # Update alpha (1/tau^2): alpha = gamma_sum / (W^T W)
        WtW = (W ** 2).sum(0).mean().item()
        if WtW > 1e-12:
            alpha = gamma_sum / (d * WtW)
        else:
            alpha = 1e6

        # Update beta (1/sigma^2) using effective degrees of freedom
        # beta = (n - gamma_sum) / RSS
        # Approximate RSS from the eigenspace
        rss = (VtXtY ** 2 * lam ** 2 / (eigvals + lam).unsqueeze(1) ** 2).sum().item() / k
        dof = max(n_samples - gamma_sum, 1.0)
        if rss > 1e-12:
            beta = dof / rss
        else:
            beta = 1e6

        lam_new = alpha / beta
        if abs(lam_new - lam) / max(abs(lam), 1e-10) < 1e-6:
            lam = lam_new
            break
        lam = lam_new

    return lam, gamma_sum, alpha, beta


def main():
    print("=" * 60)
    print("Empirical Bayes Analytical Detection Head")
    print("=" * 60, flush=True)

    manifest = json.load(open(os.path.join(CACHE_DIR, "manifest.json")))
    strides = [16, 32, 64]; H = RESOLUTION // 16
    sizes = [(H,H),(H//2,H//2),(H//4,H//4)]
    sr = [(-1,128),(128,256),(256,float("inf"))]
    locs = make_locations(sizes, strides)
    feat_dim = 768

    cls_XtX = torch.zeros(feat_dim+1, feat_dim+1, device=DEVICE)
    cls_XtY = torch.zeros(feat_dim+1, NUM_CLASSES, device=DEVICE)
    reg_XtX = torch.zeros(feat_dim+1, feat_dim+1, device=DEVICE)
    reg_XtY = torch.zeros(feat_dim+1, 4, device=DEVICE)
    ctr_XtX = torch.zeros(feat_dim+1, feat_dim+1, device=DEVICE)
    ctr_XtY = torch.zeros(feat_dim+1, 1, device=DEVICE)
    n_cls=0; n_reg=0; n_ctr=0; seen=0; n_images=20000
    t0 = time.time()

    for si in range(manifest["n_shards"]):
        if seen >= n_images: break
        shard = torch.load(os.path.join(CACHE_DIR, f"shard_{si:04d}.pt"),
                          map_location="cpu", weights_only=False)
        for item in shard:
            if seen >= n_images: break
            sp = item["spatial"].unsqueeze(0).float()
            boxes = item["boxes"]; labels = item["labels"]
            cofibers = cofiber_decompose(sp, 3)
            for sci, cof in enumerate(cofibers):
                B,C,Hc,Wc = cof.shape
                f = F.layer_norm(cof.permute(0,2,3,1).reshape(-1,C), [C]).to(DEVICE)
                ct, rt, ctrt = assign_targets(locs[sci], boxes, labels, strides[sci], sr[sci])
                pos = ct >= 0
                if not pos.any(): continue
                fp = f[pos]
                fa = torch.cat([fp, torch.ones(fp.shape[0],1,device=DEVICE)], 1)
                yc = torch.zeros(fp.shape[0], NUM_CLASSES, device=DEVICE)
                yc[torch.arange(fp.shape[0],device=DEVICE), ct[pos].to(DEVICE)] = 1.0
                cls_XtX += fa.T @ fa; cls_XtY += fa.T @ yc; n_cls += fp.shape[0]
                ltrb = rt[pos]; valid = (ltrb > 0).all(1)
                if valid.any():
                    fv = fa[valid]; yt = torch.log(ltrb[valid]).to(DEVICE)
                    reg_XtX += fv.T @ fv; reg_XtY += fv.T @ yt; n_reg += valid.sum().item()
                ctr_XtX += fa.T @ fa
                ctr_XtY += fa.T @ ctrt[pos].unsqueeze(1).to(DEVICE); n_ctr += fp.shape[0]
            seen += 1
        del shard
        if (si+1) % 5 == 0:
            print(f"  shard {si+1}: {seen} imgs, {n_cls} cls, {time.time()-t0:.0f}s", flush=True)

    print(f"\nAccumulated: {n_cls} cls, {n_reg} reg, {n_ctr} ctr", flush=True)

    # Empirical Bayes per task
    print("\nEstimating per-task lambda via empirical Bayes...", flush=True)
    t1 = time.time()
    lam_cls, gamma_cls, alpha_cls, beta_cls = empirical_bayes_lambda(cls_XtX, cls_XtY, n_cls)
    lam_reg, gamma_reg, alpha_reg, beta_reg = empirical_bayes_lambda(reg_XtX, reg_XtY, n_reg)
    lam_ctr, gamma_ctr, alpha_ctr, beta_ctr = empirical_bayes_lambda(ctr_XtX, ctr_XtY, n_ctr)
    print(f"  cls: lambda={lam_cls:.6f} (gamma={gamma_cls:.1f} effective params, alpha={alpha_cls:.4f}, beta={beta_cls:.4f})")
    print(f"  reg: lambda={lam_reg:.6f} (gamma={gamma_reg:.1f} effective params, alpha={alpha_reg:.4f}, beta={beta_reg:.4f})")
    print(f"  ctr: lambda={lam_ctr:.6f} (gamma={gamma_ctr:.1f} effective params, alpha={alpha_ctr:.4f}, beta={beta_ctr:.4f})")
    print(f"  (took {time.time()-t1:.1f}s)", flush=True)

    # Solve with EB lambdas
    I = torch.eye(feat_dim+1, device=DEVICE)
    print("\nSolving with empirical Bayes lambdas...", flush=True)
    cls_W_eb = torch.linalg.solve(cls_XtX + lam_cls * I, cls_XtY)
    reg_W_eb = torch.linalg.solve(reg_XtX + lam_reg * I, reg_XtY)
    ctr_W_eb = torch.linalg.solve(ctr_XtX + lam_ctr * I, ctr_XtY)

    # Also solve with our known-good lambda=0.1*n for comparison
    print("Solving with lambda=0.1*n (previous best)...", flush=True)
    cls_W_fix = torch.linalg.solve(cls_XtX + 0.1 * I * n_cls, cls_XtY)
    reg_W_fix = torch.linalg.solve(reg_XtX + 0.1 * I * n_reg, reg_XtY)
    ctr_W_fix = torch.linalg.solve(ctr_XtX + 0.1 * I * n_ctr, ctr_XtY)

    # Eval both
    val = torch.load(VAL_CACHE, map_location="cpu", weights_only=False)
    from pycocotools.coco import COCO
    from pycocotools.cocoeval import COCOeval
    ann_file = os.path.join(COCO_ROOT, "annotations", "instances_val2017.json")
    coco_gt = COCO(ann_file)
    cat_ids = sorted(coco_gt.getCatIds())
    idx_to_cat = {i: c for i, c in enumerate(cat_ids)}
    all_locs = torch.cat(make_locations(sizes, strides, DEVICE))

    for label, cls_W, reg_W, ctr_W in [("empirical_bayes", cls_W_eb, reg_W_eb, ctr_W_eb),
                                         ("fixed_0.1n", cls_W_fix, reg_W_fix, ctr_W_fix)]:
        print(f"\nEvaluating: {label}", flush=True)
        all_results = []
        for idx in range(len(val)):
            item = val[idx]
            spatial = item["spatial"].unsqueeze(0).float().to(DEVICE)
            img_id = int(item["img_id"]); scale = item["scale"]
            cofibers = cofiber_decompose(spatial, 3)
            cls_all, reg_all, ctr_all = [], [], []
            for cof in cofibers:
                B,C,Hc,Wc = cof.shape
                f = F.layer_norm(cof.permute(0,2,3,1).reshape(-1,C), [C])
                fa = torch.cat([f, torch.ones(f.shape[0],1,device=DEVICE)], 1)
                cls = (fa @ cls_W).sigmoid()
                reg = (fa @ reg_W).exp()
                ctr = (fa @ ctr_W).sigmoid()
                cls_all.append(cls); reg_all.append(reg); ctr_all.append(ctr.squeeze(1))
            cls_s = torch.cat(cls_all); reg_s = torch.cat(reg_all); ctr_s = torch.cat(ctr_all)
            scores = cls_s * ctr_s.unsqueeze(1)
            max_s, max_c = scores.max(1)
            topk = min(100, max_s.shape[0])
            top_s, top_i = max_s.topk(topk)
            tc = max_c[top_i]; tr = reg_s[top_i]; tl = all_locs[top_i]
            x1=(tl[:,0]-tr[:,0])/scale; y1=(tl[:,1]-tr[:,1])/scale
            x2=(tl[:,0]+tr[:,2])/scale; y2=(tl[:,1]+tr[:,3])/scale
            w=(x2-x1).clamp(min=0); h=(y2-y1).clamp(min=0)
            for i in range(topk):
                s = top_s[i].item()
                if s < 0.01: continue
                all_results.append({"image_id": img_id, "category_id": idx_to_cat[tc[i].item()],
                                    "bbox": [x1[i].item(), y1[i].item(), w[i].item(), h[i].item()],
                                    "score": s})
            if (idx+1) % 1000 == 0:
                print(f"  {idx+1}/{len(val)}", flush=True)

        if all_results:
            coco_dt = coco_gt.loadRes(all_results)
            coco_eval = COCOeval(coco_gt, coco_dt, "bbox")
            coco_eval.params.imgIds = sorted(coco_gt.getImgIds())[:len(val)]
            coco_eval.evaluate(); coco_eval.accumulate(); coco_eval.summarize()
            print(f"\n  {label}: mAP={coco_eval.stats[0]:.4f} mAP50={coco_eval.stats[1]:.4f} mAP75={coco_eval.stats[2]:.4f}")
        else:
            print(f"  {label}: no detections")

    elapsed = time.time() - t0
    print(f"\nTotal: {elapsed:.0f}s")
    print(f"EB lambdas: cls={lam_cls:.4f} reg={lam_reg:.4f} ctr={lam_ctr:.4f}")
    print(f"Fixed:      cls={0.1*n_cls:.0f} reg={0.1*n_reg:.0f} ctr={0.1*n_ctr:.0f}")


if __name__ == "__main__":
    main()