"""
Exotic regression experiments on GPU.

The classification is already at 69.6% with linear features — close to the ceiling.
The gap between analytical (1.6 mAP) and trained (8.2 mAP) is in REGRESSION.
Test whether nonlinear feature expansions help the regression solver.

Experiments:
1. Quadratic features for regression only (linear cls stays at 768)
2. Sheaf H^1 boundary features for regression only
3. Quadratic + H^1 combined
4. Full pipeline: best cls + best reg → build complete head → run actual mAP eval
"""

import json
import os
import sys
import time

import torch
import torch.nn.functional as F

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, SCRIPT_DIR)

COCO_ROOT = os.environ.get("ARENA_COCO_ROOT", "coco")
VAL_CACHE = os.environ.get("ARENA_VAL_CACHE", "val_cache/val.pt")
NUM_CLASSES = 80
DEVICE = "cuda"


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):
    locs = []
    for (h, w), s in zip(sizes, strides):
        ys = (torch.arange(h, dtype=torch.float32) + 0.5) * s
        xs = (torch.arange(w, 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_full(loc, boxes, labels, stride, sr):
    n = loc.shape[0]
    ct = torch.full((n,), -1, dtype=torch.long)
    rt = torch.zeros(n, 4)
    if boxes.numel() == 0:
        return ct, rt
    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]]
    return ct, rt


def build_val_data_with_spatial(val_path, n_images=500):
    """Build features with spatial variants on GPU."""
    val = torch.load(val_path, map_location="cpu", weights_only=False)
    from pycocotools.coco import COCO
    ann_file = os.path.join(COCO_ROOT, "annotations", "instances_val2017.json")
    coco = COCO(ann_file)
    cat_ids = sorted(coco.getCatIds())
    cat_to_idx = {c: i for i, c in enumerate(cat_ids)}
    strides = [16, 32, 64]
    H = 640 // 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)

    all_f, all_h1v, all_h1h, all_cls, all_reg = [], [], [], [], []
    for idx in range(min(n_images, len(val))):
        item = val[idx]
        spatial = item["spatial"].unsqueeze(0).float()
        img_id = item["img_id"]; scale = item["scale"]
        ann_ids = coco.getAnnIds(imgIds=int(img_id), iscrowd=False)
        anns = coco.loadAnns(ann_ids)
        boxes, labels = [], []
        for ann in anns:
            x, y, w, h = ann["bbox"]
            if w < 1 or h < 1: continue
            boxes.append([x*scale, y*scale, (x+w)*scale, (y+h)*scale])
            labels.append(cat_to_idx[ann["category_id"]])
        boxes_t = torch.tensor(boxes, dtype=torch.float32) if boxes else torch.zeros(0, 4)
        labels_t = torch.tensor(labels, dtype=torch.long) if labels else torch.zeros(0, dtype=torch.long)

        cofibers = cofiber_decompose(spatial, 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])

            # Sheaf H^1: directional boundary magnitudes
            f_4d = f.reshape(B, Hc, Wc, C).permute(0, 3, 1, 2)
            d_up = f_4d - F.pad(f_4d[:, :, 1:, :], (0, 0, 0, 1))
            d_down = f_4d - F.pad(f_4d[:, :, :-1, :], (0, 0, 1, 0))
            d_left = f_4d - F.pad(f_4d[:, :, :, 1:], (0, 1, 0, 0))
            d_right = f_4d - F.pad(f_4d[:, :, :, :-1], (1, 0, 0, 0))
            v_bound = (d_up.abs() + d_down.abs()).permute(0, 2, 3, 1).reshape(-1, C)
            h_bound = (d_left.abs() + d_right.abs()).permute(0, 2, 3, 1).reshape(-1, C)

            ct, rt = assign_targets_full(locs[sci], boxes_t, labels_t, strides[sci], sr[sci])
            all_f.append(f)
            all_h1v.append(v_bound)
            all_h1h.append(h_bound)
            all_cls.append(ct)
            all_reg.append(rt)

    features = torch.cat(all_f).to(DEVICE)
    h1v = torch.cat(all_h1v).to(DEVICE)
    h1h = torch.cat(all_h1h).to(DEVICE)
    cls_targets = torch.cat(all_cls).to(DEVICE)
    reg_targets = torch.cat(all_reg).to(DEVICE)
    return features, h1v, h1h, cls_targets, reg_targets


def solve_regression(f_pos, y_reg, f_all, reg_targets, pos_mask, lam=0.1):
    """Solve for regression weights, return quality metric."""
    valid = (y_reg > 0).all(1)
    if valid.sum() < 10:
        return 0.0
    fv = f_pos[valid]
    fa = torch.cat([fv, torch.ones(fv.shape[0], 1, device=DEVICE)], 1)
    yt = torch.log(y_reg[valid])  # log-ltrb
    fd = fv.shape[1]
    I = torch.eye(fd + 1, device=DEVICE)
    n = fv.shape[0]
    try:
        W = torch.linalg.solve(fa.T @ fa + lam * I * n, fa.T @ yt)
    except Exception:
        return 0.0

    # Quality: 1/(1+MSE) at positive locations
    pred = f_all[pos_mask] @ W[:fd] + W[fd]
    gt_ltrb = reg_targets[pos_mask]
    val2 = (gt_ltrb > 0).all(1)
    if val2.sum() < 10:
        return 0.0
    gt_log = torch.log(gt_ltrb[val2])
    pred_valid = pred[val2]
    mse = ((pred_valid - gt_log) ** 2).mean(1)
    quality = (1.0 / (1.0 + mse)).mean().item()
    return quality


def main():
    print("=" * 60)
    print("Exotic Regression Experiments (GPU)")
    print("=" * 60, flush=True)

    features, h1v, h1h, cls_targets, reg_targets = build_val_data_with_spatial(VAL_CACHE, 500)
    pos = cls_targets >= 0
    n_pos = pos.sum().item()
    f_pos = features[pos]
    reg_pos = reg_targets[pos]
    print(f"  {features.shape[0]} locations, {n_pos} positives", flush=True)

    results = []

    # Load greedy dims
    greedy_path = os.path.join(SCRIPT_DIR, "analytical_variants", "greedy_forward_gpu.json")
    greedy_dims = list(range(20))
    if os.path.isfile(greedy_path):
        with open(greedy_path) as f:
            greedy_dims = json.load(f)["selected_dims"]

    # =====================================================
    # Baseline regression: 768 raw features
    # =====================================================
    t0 = time.time()
    q = solve_regression(f_pos, reg_pos, features, reg_targets, pos)
    print(f"\n1. Baseline (768 raw): reg_quality={q:.4f} [{time.time()-t0:.2f}s]", flush=True)
    results.append({"name": "baseline_768", "reg_quality": q, "dims": 768})

    # =====================================================
    # 2. H^1 boundary features for regression
    # =====================================================
    for label, f_extra in [("h1v", h1v), ("h1h", h1h), ("h1_both", torch.cat([h1v, h1h], 1))]:
        f_combined = torch.cat([features, f_extra], 1)
        fp = f_combined[pos]
        t0 = time.time()
        q = solve_regression(fp, reg_pos, f_combined, reg_targets, pos)
        print(f"2. 768 + {label} ({f_combined.shape[1]} dims): reg_quality={q:.4f} [{time.time()-t0:.2f}s]", flush=True)
        results.append({"name": f"h1_{label}", "reg_quality": q, "dims": f_combined.shape[1]})

    # =====================================================
    # 3. Quadratic features for regression
    # =====================================================
    for K in [10, 20, 30]:
        dims = greedy_dims[:K]
        f_sub = features[:, dims]
        quads = []
        for i in range(K):
            for j in range(i, K):
                quads.append(f_sub[:, i] * f_sub[:, j])
        f_quad = torch.stack(quads, 1)
        f_exp = torch.cat([features, f_quad], 1)
        fp = f_exp[pos]
        t0 = time.time()
        q = solve_regression(fp, reg_pos, f_exp, reg_targets, pos)
        nd = f_exp.shape[1]
        print(f"3. 768 + quad_top{K} ({nd} dims): reg_quality={q:.4f} [{time.time()-t0:.2f}s]", flush=True)
        results.append({"name": f"quad_top{K}", "reg_quality": q, "dims": nd})

    # =====================================================
    # 4. H^1 + quadratic combined
    # =====================================================
    dims = greedy_dims[:20]
    f_sub = features[:, dims]
    quads = []
    for i in range(20):
        for j in range(i, 20):
            quads.append(f_sub[:, i] * f_sub[:, j])
    f_quad = torch.stack(quads, 1)
    f_all = torch.cat([features, h1v, h1h, f_quad], 1)
    fp = f_all[pos]
    t0 = time.time()
    q = solve_regression(fp, reg_pos, f_all, reg_targets, pos)
    print(f"4. 768 + H1 + quad_top20 ({f_all.shape[1]} dims): reg_quality={q:.4f} [{time.time()-t0:.2f}s]", flush=True)
    results.append({"name": "h1_quad_combined", "reg_quality": q, "dims": f_all.shape[1]})

    # =====================================================
    # 5. RFF for regression
    # =====================================================
    sub = features[:5000]
    dists = torch.cdist(sub[:500], sub[:500])
    sigma = dists.median().item()
    if sigma < 1e-6: sigma = 1.0

    for K_rff in [100, 500]:
        torch.manual_seed(42)
        W_rff = torch.randn(768, K_rff, device=DEVICE) / sigma
        b_rff = torch.rand(K_rff, device=DEVICE) * 2 * 3.14159
        rff = (2.0 / K_rff) ** 0.5 * torch.cos(features @ W_rff + b_rff)
        f_combined = torch.cat([features, rff], 1)
        fp = f_combined[pos]
        t0 = time.time()
        q = solve_regression(fp, reg_pos, f_combined, reg_targets, pos)
        print(f"5. 768 + {K_rff} RFF ({f_combined.shape[1]} dims): reg_quality={q:.4f} [{time.time()-t0:.2f}s]", flush=True)
        results.append({"name": f"rff_{K_rff}_reg", "reg_quality": q, "dims": f_combined.shape[1]})

    # =====================================================
    # Summary
    # =====================================================
    print(f"\n{'='*60}")
    print("Ranked by regression quality:")
    for r in sorted(results, key=lambda x: -x["reg_quality"]):
        print(f"  {r['name']:25s}: reg_quality={r['reg_quality']:.4f}  dims={r['dims']}")

    out = os.path.join(SCRIPT_DIR, "analytical_variants", "exotic_reg_gpu.json")
    with open(out, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nSaved: {out}")


if __name__ == "__main__":
    main()