"""
Multi-Domain Detection Arena — train and evaluate heads across COCO + RF20-VL.

Caches frozen backbone features per domain, trains each candidate head on
each domain's training split, evaluates precision/recall on the validation
split. Produces a single leaderboard across all domains.

Usage:
    python multi_domain_arena.py --steps 2000 --batch 1 --candidates S,R,A
    python multi_domain_arena.py --cache-only
    python multi_domain_arena.py --domains aquarium-combined,aerial-airport
"""

import argparse
import json
import math
import os
import sys
import time
from typing import Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torch import Tensor
from torchvision.ops import nms
from torchvision.transforms import v2

# Import all candidates from the arena
sys.path.insert(0, os.path.dirname(__file__))
from losses.fcos import assign_targets, focal_loss, fcos_loss, FPN_STRIDES, SIZE_RANGES
from losses.centernet import centernet_targets, centernet_loss
from utils.decode import make_locations, decode_fcos, decode_centernet
from utils.cache import letterbox, HOOK_BLOCKS, N_PREFIX
from heads import REGISTRY

BACKBONE_REPO = os.environ.get("ARENA_BACKBONE_REPO", "./backbone")
BACKBONE_WEIGHTS = os.environ.get("ARENA_BACKBONE_WEIGHTS", "./backbone/weights.pt")
BACKBONE_HUB_ENTRY = os.environ.get("ARENA_BACKBONE_ENTRY", "eupe_vitb16")
CACHE_DIR = os.environ.get("ARENA_CACHE_DIR", "./arena_cache")
RF20_ROOT = os.environ.get("ARENA_RF20_ROOT", "./rf100vl/rf20")
COCO_ROOT = os.environ.get("ARENA_COCO_ROOT", "./coco")
RESOLUTION = 640
N_TRAIN_PER_DOMAIN = 300
N_DIAG_PER_DOMAIN = 1000

if BACKBONE_REPO not in sys.path:
    sys.path.insert(0, BACKBONE_REPO)


# ---------------------------------------------------------------------------
# Domain dataset loading (COCO format, variable num_classes)
# ---------------------------------------------------------------------------

def load_domain_dataset(img_dir, ann_file, n_images, min_anns=0):
    """Load images + annotations from a COCO-format dataset."""
    with open(ann_file) as f:
        coco = json.load(f)

    cat_ids = sorted(set(c["id"] for c in coco["categories"]))
    cat_to_contig = {cat: i for i, cat in enumerate(cat_ids)}
    cat_names = {c["id"]: c["name"] for c in coco["categories"]}
    num_classes = len(cat_ids)

    id_to_anns = {}
    for a in coco["annotations"]:
        if a.get("iscrowd", 0):
            continue
        cat = a["category_id"]
        if cat not in cat_to_contig:
            continue
        id_to_anns.setdefault(a["image_id"], []).append(a)

    id_to_info = {img["id"]: img for img in coco["images"]}
    candidates = [(iid, anns) for iid, anns in id_to_anns.items() if len(anns) >= min_anns]

    import random
    random.seed(42)
    random.shuffle(candidates)

    items = []
    for iid, anns in candidates[:n_images]:
        info = id_to_info[iid]
        path = os.path.join(img_dir, info["file_name"])
        if not os.path.isfile(path):
            continue
        boxes, labels = [], []
        for a in anns:
            x, y, w, h = a["bbox"]
            if w < 1 or h < 1:
                continue
            boxes.append([x, y, x + w, y + h])
            labels.append(cat_to_contig[a["category_id"]])
        if boxes:
            items.append({"path": path, "boxes": boxes, "labels": labels,
                          "width": info["width"], "height": info["height"]})

    return items, num_classes, [cat_names.get(c, f"cls_{c}") for c in cat_ids]


def discover_domains(rf20_root, coco_root):
    """Find all available domains (RF20-VL + COCO)."""
    domains = []

    # COCO
    coco_ann = os.path.join(coco_root, "annotations", "instances_train2017.json")
    coco_val_ann = os.path.join(coco_root, "annotations", "instances_val2017.json")
    if os.path.isfile(coco_ann):
        domains.append({
            "name": "coco",
            "train_img": os.path.join(coco_root, "train2017"),
            "train_ann": coco_ann,
            "val_img": os.path.join(coco_root, "val2017"),
            "val_ann": coco_val_ann,
            "category": "Standard",
        })

    # RF20-VL
    if os.path.isdir(rf20_root):
        for d in sorted(os.listdir(rf20_root)):
            dp = os.path.join(rf20_root, d)
            train_ann = os.path.join(dp, "train", "_annotations.coco.json")
            val_ann = os.path.join(dp, "valid", "_annotations.coco.json")
            if os.path.isfile(train_ann) and os.path.isfile(val_ann):
                domains.append({
                    "name": d,
                    "train_img": os.path.join(dp, "train"),
                    "train_ann": train_ann,
                    "val_img": os.path.join(dp, "valid"),
                    "val_ann": val_ann,
                    "category": "RF20-VL",
                })

    return domains


# ---------------------------------------------------------------------------
# Feature caching per domain
# ---------------------------------------------------------------------------

def cache_domain_features(backbone, domain, cache_dir):
    """Cache features for one domain's train and val splits."""
    domain_cache = os.path.join(cache_dir, "domains", domain["name"])
    train_path = os.path.join(domain_cache, "train.pt")
    val_path = os.path.join(domain_cache, "val.pt")

    if os.path.isfile(train_path) and os.path.isfile(val_path):
        return True

    os.makedirs(domain_cache, exist_ok=True)
    normalize = v2.Compose([
        v2.ToImage(), v2.ToDtype(torch.float32, scale=True),
        v2.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ])

    hooks, intermediates = [], {}
    for idx in HOOK_BLOCKS:
        def _hook(block_idx):
            def fn(mod, inp, out):
                intermediates[block_idx] = (out[0] if isinstance(out, list) else out).detach()
            return fn
        hooks.append(backbone.blocks[idx].register_forward_hook(_hook(idx)))

    for split, n, tag in [("train", N_TRAIN_PER_DOMAIN, "train"), ("val", N_DIAG_PER_DOMAIN, "val")]:
        out_path = os.path.join(domain_cache, f"{tag}.pt")
        if os.path.isfile(out_path):
            continue

        if tag == "train":
            items, num_classes, class_names = load_domain_dataset(
                domain["train_img"], domain["train_ann"], n, min_anns=1)
        else:
            items, num_classes, class_names = load_domain_dataset(
                domain["val_img"], domain["val_ann"], n, min_anns=1)

        cached = []
        for item in items:
            img = Image.open(item["path"]).convert("RGB")
            canvas, scale = letterbox(img, RESOLUTION)
            x = normalize(canvas).unsqueeze(0).cuda()

            intermediates.clear()
            with torch.no_grad():
                with torch.autocast("cuda", dtype=torch.bfloat16):
                    out = backbone.forward_features(x)
            patches = out["x_norm_patchtokens"].float()
            B, N, D = patches.shape
            h = w = int(N ** 0.5)
            spatial = patches[0].permute(1, 0).reshape(D, h, w).half().cpu()
            inter = [intermediates[idx][0].half().cpu() for idx in HOOK_BLOCKS]

            boxes = torch.tensor(item["boxes"], dtype=torch.float32) * scale
            labels = torch.tensor(item["labels"], dtype=torch.long)

            # Precompute FCOS targets for single-level (stride 16) and multi-level (FPN)
            single_locs = make_locations([(h, w)], [16], boxes.device)
            single_cls, single_reg, single_ctr = assign_targets(
                single_locs, boxes, labels, [16], [(-1, float("inf"))])

            # Multi-level: cofiber scales [16, 32, 64]
            cofiber_sizes = [(h, w), (h // 2, w // 2), (h // 4, w // 4)]
            cofiber_strides = [16, 32, 64]
            cofiber_ranges = [(-1, 128), (128, 256), (256, float("inf"))]
            cofiber_locs = make_locations(cofiber_sizes, cofiber_strides, boxes.device)
            cofiber_cls, cofiber_reg, cofiber_ctr = assign_targets(
                cofiber_locs, boxes, labels, cofiber_strides, cofiber_ranges)

            cached.append({
                "spatial": spatial, "intermediates": inter,
                "boxes": boxes, "labels": labels,
                "targets_single": (single_cls, single_reg, single_ctr),
                "targets_cofiber": (cofiber_cls, cofiber_reg, cofiber_ctr),
            })

        torch.save({"items": cached, "num_classes": num_classes, "class_names": class_names}, out_path)
        print(f"    {domain['name']}/{tag}: {len(cached)} images, {num_classes} classes", flush=True)

    for h in hooks:
        h.remove()
    return True


# ---------------------------------------------------------------------------
# Per-domain training and eval (adapted from detection_arena.run_candidate)
# ---------------------------------------------------------------------------

def make_candidate(key, num_classes):
    """Instantiate a candidate with the correct num_classes for this domain."""
    # Most candidates hardcode NUM_CLASSES=80. For multi-domain we need to
    # override. The simplest approach: instantiate then replace the cls head.
    # For candidates with fundamentally different architectures, we handle
    # them individually.
    # Letter keys for backward compat with old scripts
    LETTER_TO_NAME = {
        "A": "baseline_fcos", "B": "slim_fcos", "C": "centernet",
        "D": "hook_fcos", "E": "sparse_query", "F": "prototype",
        "G": "depth_fusion", "H": "wavelet", "I": "adaptive_query",
        "J": "patch_assembly", "K": "mutual_attention", "M": "feature_graph",
        "N": "scale_classify", "O": "relational_corners", "P": "cascade_pool",
        "Q": "cofiber_linear", "R": "threshold_prototype", "S": "cofiber_threshold",
        "S1": "cofiber_centernet", "S2": "cofiber_5scale", "S3": "cofiber_adaptive",
    }
    name = LETTER_TO_NAME.get(key, key)
    if name not in REGISTRY:
        return None
    return REGISTRY[name]()


def run_on_domain(candidate_key, domain_name, cache_dir, steps=2000, batch_size=1, lr=1e-3, seed=42):
    """Train and eval one candidate on one domain. Returns result dict."""
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

    domain_cache = os.path.join(cache_dir, "domains", domain_name)
    train_data_raw = torch.load(os.path.join(domain_cache, "train.pt"), map_location="cpu", weights_only=False)
    val_data_raw = torch.load(os.path.join(domain_cache, "val.pt"), map_location="cpu", weights_only=False)

    train_items = train_data_raw["items"]
    val_items = val_data_raw["items"]
    num_classes = train_data_raw["num_classes"]

    if len(train_items) == 0 or len(val_items) == 0:
        return {"name": candidate_key, "domain": domain_name, "num_classes": num_classes,
                "loss_end": float("nan"), "precision": 0.0, "recall": 0.0,
                "tp": 0, "fp": 0, "n_gt": 0, "skip": True}

    candidate = make_candidate(candidate_key, num_classes)
    if candidate is None:
        return None
    candidate = candidate.cuda()

    optimizer = torch.optim.AdamW(candidate.parameters(), lr=lr, weight_decay=1e-4)
    candidate.train()

    sample_sp = train_items[0]["spatial"].unsqueeze(0).float().cuda()
    sample_inter = [t.unsqueeze(0).float().cuda() for t in train_items[0]["intermediates"]] if candidate.needs_intermediates else None
    with torch.no_grad():
        locs = candidate.get_locs(sample_sp)

    n = len(train_items)
    losses = []
    t0 = time.time()

    # Determine which precomputed target key to use based on head type
    uses_cofiber = hasattr(candidate, 'n_scales') and candidate.n_scales > 1
    target_key = "targets_cofiber" if uses_cofiber else "targets_single"

    # Check if precomputed targets exist
    # Always use the full loss (cls + box + centerness) for correct training.
    # The precomputed cls-only path is disabled — it trains classification
    # without box regression, producing models that cannot detect.
    has_precomputed = False

    for step in range(steps):
        indices = torch.randint(0, n, (min(batch_size, n),))
        spatial_b = torch.stack([train_items[i]["spatial"] for i in indices]).float().cuda()

        if candidate.needs_intermediates:
            inter_b = [torch.stack([train_items[i]["intermediates"][j] for i in indices]).float().cuda()
                       for j in range(4)]
        else:
            inter_b = None

        preds = candidate(spatial_b, inter_b)

        if has_precomputed:
            # Fast path: classification-only focal loss on precomputed targets.
            # Skip GIoU and centerness — they're expensive and unnecessary for
            # screening. Class prediction quality is the signal we need.
            cls_levels = preds[0]
            n_levels = len(cls_levels)
            B_actual = cls_levels[0].shape[0]
            num_cls = cls_levels[0].shape[1]
            device = cls_levels[0].device

            # Stack all predictions and targets across batch and levels
            all_pred, all_tgt = [], []
            for li in range(n_levels):
                pred_c = cls_levels[li].permute(0, 2, 3, 1).reshape(B_actual, -1, num_cls)
                for bi_local, bi_data in enumerate(indices):
                    tgt_cls = train_items[bi_data.item()][target_key][0][li]
                    all_pred.append(pred_c[bi_local])
                    all_tgt.append(tgt_cls)

            all_pred = torch.cat(all_pred, 0).to(device)
            all_tgt = torch.cat(all_tgt, 0).to(device)
            pos = all_tgt >= 0
            npos = max(pos.sum().item(), 1)
            oh = torch.zeros(len(all_tgt), num_cls, device=device)
            if pos.any():
                oh[pos, all_tgt[pos].clamp(0, num_cls - 1)] = 1.0
            loss = focal_loss(all_pred, oh) / npos
        else:
            # Slow path: compute targets on the fly
            boxes_b = [train_items[i]["boxes"].cuda() for i in indices]
            labels_b = [train_items[i]["labels"].cuda() for i in indices]
            loss = candidate.loss(preds, locs, boxes_b, labels_b)

        if torch.isnan(loss) or torch.isinf(loss):
            losses.append(float("nan"))
            optimizer.zero_grad()
            continue

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(candidate.parameters(), 5.0)
        optimizer.step()
        losses.append(loss.item())

        if (step + 1) % 500 == 0:
            avg = np.mean(losses[-100:]) if len(losses) >= 100 else np.mean(losses)
            print(f"      step {step+1}/{steps}  loss={avg:.4f}  ({time.time()-t0:.1f}s)", flush=True)

    train_time = time.time() - t0

    # Eval on val split
    candidate.eval()
    total_tp, total_fp, total_gt = 0, 0, 0
    with torch.no_grad():
        for item in val_items:
            sp = item["spatial"].unsqueeze(0).float().cuda()
            inter = [t.unsqueeze(0).float().cuda() for t in item["intermediates"]] if candidate.needs_intermediates else None
            preds = candidate(sp, inter)
            dets = candidate.decode(preds, locs, score_thresh=0.3)
            det_boxes = dets[0]["boxes"]
            det_labels = dets[0]["labels"]

            gt_boxes = item["boxes"].cuda()
            gt_labels = item["labels"].cuda()
            n_gt = len(gt_labels)
            total_gt += n_gt

            if len(det_boxes) > 0 and n_gt > 0:
                x1 = torch.maximum(det_boxes[:, None, 0], gt_boxes[None, :, 0])
                y1 = torch.maximum(det_boxes[:, None, 1], gt_boxes[None, :, 1])
                x2 = torch.minimum(det_boxes[:, None, 2], gt_boxes[None, :, 2])
                y2 = torch.minimum(det_boxes[:, None, 3], gt_boxes[None, :, 3])
                inter_area = (x2 - x1).clamp(min=0) * (y2 - y1).clamp(min=0)
                det_area = (det_boxes[:, 2] - det_boxes[:, 0]) * (det_boxes[:, 3] - det_boxes[:, 1])
                gt_area = (gt_boxes[:, 2] - gt_boxes[:, 0]) * (gt_boxes[:, 3] - gt_boxes[:, 1])
                union = det_area[:, None] + gt_area[None, :] - inter_area
                iou = inter_area / union.clamp(min=1e-6)

                matched_gt = set()
                for di in range(len(det_boxes)):
                    best_iou, best_gi = iou[di].max(0)
                    gi = best_gi.item()
                    if best_iou.item() >= 0.5 and gi not in matched_gt and det_labels[di] == gt_labels[gi]:
                        total_tp += 1
                        matched_gt.add(gi)
                    else:
                        total_fp += 1
            else:
                total_fp += len(det_boxes)

    precision = total_tp / max(total_tp + total_fp, 1)
    recall = total_tp / max(total_gt, 1)

    del candidate
    torch.cuda.empty_cache()

    return {
        "name": candidate_key, "domain": domain_name, "num_classes": num_classes,
        "loss_end": np.mean(losses[-10:]) if losses else float("nan"),
        "precision": precision, "recall": recall,
        "tp": total_tp, "fp": total_fp, "n_gt": total_gt,
        "train_time_s": train_time, "skip": False,
    }


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--steps", type=int, default=2000)
    parser.add_argument("--batch", type=int, default=16)
    parser.add_argument("--cache-only", action="store_true")
    parser.add_argument("--candidates", type=str, default="S,R,A,C")
    parser.add_argument("--domains", type=str, default="all")
    args = parser.parse_args()

    print("=" * 80)
    print("Multi-Domain Detection Arena")
    print("=" * 80)

    # Discover domains
    domains = discover_domains(RF20_ROOT, COCO_ROOT)
    if args.domains != "all":
        selected = set(d.strip() for d in args.domains.split(","))
        domains = [d for d in domains if d["name"] in selected]
    print(f"\n  {len(domains)} domains available")
    for d in domains:
        print(f"    {d['name']} ({d['category']})")

    # Load backbone and cache features
    print("\n  Loading backbone...")
    backbone = torch.hub.load(BACKBONE_REPO, BACKBONE_HUB_ENTRY, source="local", weights=BACKBONE_WEIGHTS)
    backbone = backbone.cuda().eval()
    for p in backbone.parameters():
        p.requires_grad = False

    print("\n  Caching features per domain...")
    for d in domains:
        cache_domain_features(backbone, d, CACHE_DIR)
    del backbone
    torch.cuda.empty_cache()

    if args.cache_only:
        print("\n  Cache complete.")
        return

    # Run candidates across domains
    candidate_keys = [c.strip().upper() for c in args.candidates.split(",")]
    total_runs = len(candidate_keys) * len(domains)
    print(f"\n  Running {len(candidate_keys)} candidates x {len(domains)} domains = {total_runs} runs")
    print(f"  Steps: {args.steps}, batch: {args.batch}, seed: 42", flush=True)

    all_results = []
    run_count = 0
    sweep_t0 = time.time()
    errors = []

    for ci, ck in enumerate(candidate_keys):
        cand_t0 = time.time()
        cand_results = []
        print(f"\n  --- Candidate {ck} ({ci+1}/{len(candidate_keys)}) ---", flush=True)

        for di, d in enumerate(domains):
            run_count += 1
            try:
                result = run_on_domain(ck, d["name"], CACHE_DIR, steps=args.steps, batch_size=args.batch)
                if result is None:
                    continue
                all_results.append(result)
                cand_results.append(result)

                elapsed = time.time() - sweep_t0
                rate = run_count / elapsed
                remaining = (total_runs - run_count) / rate if rate > 0 else 0
                eta_m = remaining / 60

                if result.get("skip"):
                    print(f"  [{run_count}/{total_runs}] {ck} x {d['name']:<25} SKIP  (ETA {eta_m:.0f}m)", flush=True)
                else:
                    print(f"  [{run_count}/{total_runs}] {ck} x {d['name']:<25} "
                          f"loss={result['loss_end']:.3f} prec={result['precision']:.3f} "
                          f"rec={result['recall']:.3f} tp={result['tp']}  "
                          f"({result['train_time_s']:.0f}s)  ETA {eta_m:.0f}m", flush=True)

            except Exception as e:
                run_count += 1
                err_msg = f"{ck} x {d['name']}: {type(e).__name__}: {e}"
                errors.append(err_msg)
                print(f"  [{run_count}/{total_runs}] {ck} x {d['name']:<25} ERROR: {e}", flush=True)
                torch.cuda.empty_cache()

        # Per-candidate summary
        valid = [r for r in cand_results if not r.get("skip")]
        if valid:
            avg_prec = np.mean([r["precision"] for r in valid])
            avg_rec = np.mean([r["recall"] for r in valid])
            total_tp = sum(r["tp"] for r in valid)
            cand_time = time.time() - cand_t0
            print(f"  >>> {ck}: {len(valid)} domains, avg_prec={avg_prec:.4f}, "
                  f"avg_rec={avg_rec:.4f}, total_tp={total_tp}, time={cand_time:.0f}s", flush=True)

    # Final summary
    print("\n" + "=" * 80)
    print("CROSS-DOMAIN SUMMARY")
    print("=" * 80)
    print(f"{'Candidate':<20} {'Domains':>7} {'AvgPrec':>8} {'AvgRec':>8} {'TotalTP':>8} {'AvgLoss':>8}")
    print("-" * 70)
    for ck in candidate_keys:
        ck_results = [r for r in all_results if r["name"] == ck and not r.get("skip")]
        if not ck_results:
            continue
        avg_prec = np.mean([r["precision"] for r in ck_results])
        avg_rec = np.mean([r["recall"] for r in ck_results])
        total_tp = sum(r["tp"] for r in ck_results)
        avg_loss = np.mean([r["loss_end"] for r in ck_results])
        print(f"{ck:<20} {len(ck_results):>7} {avg_prec:>8.4f} {avg_rec:>8.4f} {total_tp:>8} {avg_loss:>8.3f}")

    if errors:
        print(f"\n  {len(errors)} errors:")
        for e in errors:
            print(f"    {e}")

    total_time = time.time() - sweep_t0
    print(f"\n  Total time: {total_time/60:.1f} minutes ({total_time/3600:.1f} hours)")

    # Save
    out = os.path.join(CACHE_DIR, "multi_domain_results.json")
    with open(out, "w") as f:
        json.dump(all_results, f, indent=2, default=str)
    print(f"\nSaved to {out}")


if __name__ == "__main__":
    main()