"""
GCV-optimal analytical detection head.

Computes the generalized cross-validation optimal lambda for each task
(classification, regression, centerness) independently via SVD of the
accumulated sufficient statistics. No grid search — closed-form.

Then builds the head with per-task optimal regularization and evals.
"""

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

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
STATS_DIR = os.path.join(SCRIPT_DIR, "analytical_stats_cache")
COCO_ROOT = os.environ.get("ARENA_COCO_ROOT")
VAL_CACHE = os.environ.get("ARENA_VAL_CACHE")
CACHE_DIR = os.environ.get("ARENA_CACHE_DIR")
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 gcv_optimal_lambda(XtX, XtY, n_samples, lambdas=None):
    """Find GCV-optimal lambda via SVD.

    For ridge regression: Y = X W + noise
    Hat matrix: H(λ) = X (X^T X + λI)^{-1} X^T
    GCV(λ) = (1/n) ||Y - H(λ)Y||² / (1 - tr(H(λ))/n)²

    Using SVD of the augmented feature matrix: X = U S V^T
    tr(H(λ)) = Σ s_i² / (s_i² + λ)
    ||Y - H(λ)Y||² = Σ (λ u_i^T Y / (s_i² + λ))²

    We work with XtX = V S² V^T directly.
    """
    # Eigendecompose XtX (symmetric positive semi-definite)
    eigvals, eigvecs = torch.linalg.eigh(XtX)
    eigvals = eigvals.clamp(min=0)  # numerical stability

    # Project targets: V^T X^T Y = V^T (XtX XtX^{-1} XtY) ...
    # Actually we need U^T Y. From XtX = V S² V^T and XtY = V S U^T Y
    # So V^T XtY = S U^T Y, thus U^T Y = S^{-1} V^T XtY
    VtXtY = eigvecs.T @ XtY  # (d+1, k)
    # S = sqrt(eigvals)
    S = eigvals.sqrt().clamp(min=1e-10)
    UtY = VtXtY / S.unsqueeze(1)  # (d+1, k)

    if lambdas is None:
        # Log-spaced search
        lam_min = eigvals[eigvals > 0].min().item() * 0.001
        lam_max = eigvals.max().item() * 10
        lambdas = torch.logspace(
            max(-8, torch.log10(torch.tensor(lam_min)).item()),
            min(4, torch.log10(torch.tensor(lam_max)).item()),
            200, device=XtX.device)

    d = XtX.shape[0]
    best_lam = lambdas[0].item()
    best_gcv = float("inf")

    for lam in lambdas:
        # tr(H) = Σ s_i² / (s_i² + λ)
        leverage = eigvals / (eigvals + lam)
        tr_H = leverage.sum()

        # Residual: ||Y - HY||² = Σ_i (λ/(s_i²+λ))² ||u_i^T Y||²
        shrinkage = lam / (eigvals + lam)
        res_sq = (shrinkage.unsqueeze(1) ** 2 * UtY ** 2).sum()

        # GCV
        gcv = (res_sq / n_samples) / ((1 - tr_H / n_samples) ** 2 + 1e-12)

        if gcv.item() < best_gcv:
            best_gcv = gcv.item()
            best_lam = lam.item()

    return best_lam, best_gcv


def main():
    print("=" * 60)
    print("GCV-Optimal Analytical Detection Head")
    print("=" * 60, flush=True)

    # Accumulate on GPU
    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
    n_images = 20000; seen = 0
    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, {n_reg} reg, {time.time()-t0:.0f}s", flush=True)

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

    # GCV optimal lambda per task
    print("\nFinding GCV-optimal lambda...", flush=True)
    t1 = time.time()
    # Normalize by n so GCV searches over the actual regularization strength
    # Our solve uses (XtX + lam * I * n), so GCV should find lam such that lam*n is optimal
    # Equivalently: search on (XtX/n + lam * I) and report lam
    lam_cls, gcv_cls = gcv_optimal_lambda(cls_XtX / n_cls, cls_XtY / n_cls, n_cls)
    lam_reg, gcv_reg = gcv_optimal_lambda(reg_XtX / n_reg, reg_XtY / n_reg, n_reg)
    lam_ctr, gcv_ctr = gcv_optimal_lambda(ctr_XtX / n_ctr, ctr_XtY / n_ctr, n_ctr)
    print(f"  cls: lambda={lam_cls:.6f} (GCV={gcv_cls:.6f})")
    print(f"  reg: lambda={lam_reg:.6f} (GCV={gcv_reg:.6f})")
    print(f"  ctr: lambda={lam_ctr:.6f} (GCV={gcv_ctr:.6f})")
    print(f"  (took {time.time()-t1:.1f}s)", flush=True)

    # Compare: solve with GCV lambda vs fixed lambda=0.1
    print("\nSolving with GCV-optimal lambda...", flush=True)
    I = torch.eye(feat_dim+1, device=DEVICE)
    cls_W_gcv = torch.linalg.solve(cls_XtX + lam_cls * I * n_cls, cls_XtY)
    reg_W_gcv = torch.linalg.solve(reg_XtX + lam_reg * I * n_reg, reg_XtY)
    ctr_W_gcv = torch.linalg.solve(ctr_XtX + lam_ctr * I * n_ctr, ctr_XtY)

    print("Solving with fixed lambda=0.1...", 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 on COCO val
    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 [("gcv_optimal", cls_W_gcv, reg_W_gcv, ctr_W_gcv),
                                         ("fixed_0.1", 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()
            mAP = coco_eval.stats[0]
            mAP50 = coco_eval.stats[1]
            mAP75 = coco_eval.stats[2]
            print(f"\n  {label}: mAP={mAP:.4f} mAP50={mAP50:.4f} mAP75={mAP75:.4f}")
        else:
            print(f"  {label}: no detections")

    elapsed = time.time() - t0
    print(f"\nTotal: {elapsed:.0f}s")
    print(f"\nGCV-optimal lambdas: cls={lam_cls:.6f} reg={lam_reg:.6f} ctr={lam_ctr:.6f}")


if __name__ == "__main__":
    main()