# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch",
#     "transformers>=4.40",
#     "datasets>=2.20",
#     "pillow",
#     "accelerate",
#     "pycocotools",
#     "huggingface_hub",
#     "trackio",
#     "timm",
#     "scipy",
# ]
# ///
import os
import sys
import copy
import argparse
import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset

from datasets import load_dataset
from transformers import DetrImageProcessor, DetrForObjectDetection

ID2LABEL = {
    0: "Photograph",
    1: "Illustration",
    2: "Map",
    3: "Comics/Cartoon",
    4: "Editorial Cartoon",
    5: "Headline",
    6: "Advertisement",
}
LABEL2ID = {v: k for k, v in ID2LABEL.items()}
CLASSES = [ID2LABEL[i] for i in sorted(ID2LABEL)]

MODEL_ID = "facebook/detr-resnet-50"  # Apache-2.0
REPO_ID = "harness-race/opencode-r2"

# ----- trackio (best effort) -----
def setup_trackio():
    if os.environ.get("OPENCODE_TRACKIO", "1") == "0":
        return (lambda **kw: None)
    try:
        import trackio
        r = trackio.init(project="opencode-r2", name=os.environ.get("JOB_NAME", "finetune"), private=True)

        def tlog(**kw):
            try:
                for k, v in kw.items():
                    trackio.log(f"metric/{k}", v)
            except Exception:
                pass
        return tlog
    except Exception as e:
        print("[trackio] unavailable:", e)
        return (lambda **kw: None)


# ---------- dataset ----------
class DetrDataset(Dataset):
    def __init__(self, hf_ds, processor, split):
        self.ds = hf_ds
        self.processor = processor
        self.split = split

    def __len__(self):
        return len(self.ds)

    def __getitem__(self, idx):
        item = self.ds[idx]
        image = item["image"].convert("RGB")
        objs = item["objects"]  # list of per-object dicts
        bboxes_xywh = []
        cat_ids = []
        for o in objs:
            box = o["bbox"]
            crowd = bool(o["iscrowd"]) if o.get("iscrowd") is not None else False
            if box[2] <= 0 or box[3] <= 0:
                continue
            if crowd and self.split == "train":
                continue
            bboxes_xywh.append(box)
            cat_ids.append(o["category_id"])  # int label 0..6
        annotations = [
            {"area": b[2] * b[3], "bbox": list(b), "category_id": c}
            for b, c in zip(bboxes_xywh, cat_ids)
        ]
        anno = {"image_id": int(item["image_id"]), "annotations": annotations}
        encoding = self.processor(
            images=image,
            annotations=anno,
            return_tensors="pt",
        )
        if "labels" in encoding and isinstance(encoding["labels"], list) and len(encoding["labels"]) == 1:
            encoding["labels"] = encoding["labels"][0]
        for k, v in encoding.items():
            if k == "labels":
                continue
            if isinstance(v, torch.Tensor) and len(v.shape) > 0:
                try:
                    encoding[k] = v.squeeze(0)
                except Exception:
                    pass
        encoding["image_id"] = int(item["image_id"])
        encoding["image_size"] = [int(item["height"]), int(item["width"])]
        return encoding


def collate_fn(batch):
    H = max(b["pixel_values"].shape[-2] for b in batch)
    W = max(b["pixel_values"].shape[-1] for b in batch)
    pixel_values = []
    pixel_mask = []
    for b in batch:
        im = torch.as_tensor(b["pixel_values"])
        h, w = im.shape[-2:]
        if (h, w) != (H, W):
            im = torch.nn.functional.pad(im, (0, W - w, 0, H - h), value=0.0)
        m = torch.zeros((H, W), dtype=torch.int64)
        m[:h, :w] = 1
        pixel_values.append(im)
        pixel_mask.append(m)
    pixel_values = torch.stack(pixel_values)
    pixel_mask = torch.stack(pixel_mask)
    labels = []
    for b in batch:
        lb = None
        if "labels" in b and b["labels"] is not None and "boxes" in b["labels"]:
            lab = b["labels"]
            lb = {
                "class_labels": lab["class_labels"].clone() if isinstance(lab["class_labels"], torch.Tensor) else torch.tensor(lab["class_labels"], dtype=torch.long),
                "boxes": lab["boxes"].clone() if isinstance(lab["boxes"], torch.Tensor) else torch.tensor(lab["boxes"], dtype=torch.float32),
            }
            if lb["boxes"].numel() == 0:
                lb["boxes"] = torch.zeros((0, 4), dtype=torch.float32)
        else:
            lb = {"class_labels": torch.zeros((0,), dtype=torch.long),
                  "boxes": torch.zeros((0, 4), dtype=torch.float32)}
        lb["image_id"] = b["image_id"]
        lb["image_size"] = b["image_size"]
        labels.append(lb)
    return {"pixel_values": pixel_values, "pixel_mask": pixel_mask, "labels": labels}


# ---------- coco evaluation ----------
def to_coco(preds, gts, all_image_ids):
    """preds: list of {image_id, score, label, box_xyxy(pixels)}
       gts: list of {image_id, category_id, bbox_xywh, area, ann_id}
    """
    cat_id_map = {i: i + 1 for i in range(7)}  # 0..6 -> 1..7
    all_image_ids = list(dict.fromkeys(all_image_ids))
    data = {
        "images": [{"id": int(im)} for im in all_image_ids],
        "categories": [{"id": i + 1, "name": CLASSES[i]} for i in range(7)],
        "annotations": [
            {"id": g["ann_id"], "image_id": g["image_id"], "category_id": cat_id_map[g["category_id"]],
             "bbox": g["bbox_xywh"], "area": g["area"], "iscrowd": 0}
            for g in gts
        ],
    }
    from pycocotools.coco import COCO
    from pycocotools.cocoeval import COCOeval
    coco_gt = COCO()
    coco_gt.dataset = data
    coco_gt.createIndex()
    res = []
    for p in preds:
        res.append({
            "image_id": p["image_id"], "category_id": cat_id_map[p["label"]],
            "bbox": [p["box_xyxy"][0], p["box_xyxy"][1],
                     p["box_xyxy"][2] - p["box_xyxy"][0], p["box_xyxy"][3] - p["box_xyxy"][1]],
            "score": float(p["score"]),
        })
    if not res:
        return {"mAP": 0.0, "AP50": 0.0, "AP75": 0.0, "AR1": 0.0, "AR10": 0.0, "AR100": 0.0}
    coco_dt = coco_gt.loadRes(res)
    e = COCOeval(coco_gt, coco_dt, "bbox")
    e.evaluate()
    e.accumulate()
    e.summarize()
    s = e.stats
    out = {"mAP": float(s[0]), "AP50": float(s[1]), "AP75": float(s[2]),
           "AR1": float(s[6]), "AR10": float(s[7]), "AR100": float(s[8])}
    prec = e.eval["precision"]  # (T=10 IoU, R=101 rec, K cat, A=4 area, M=3 maxDet)
    per = {}
    for i in range(7):
        p = prec[:, :, i, 0, 2].flatten()  # all IoU, all rec, class i, area=all, maxDet=100
        p = p[p > -1]
        per[CLASSES[i]] = float(p.mean()) if p.size > 0 else 0.0
    out["per_class_mAP"] = per
    p50 = {}
    for i in range(7):
        p = prec[0, :, i, 0, 2].flatten()  # IoU=0.5
        p = p[p > -1]
        p50[CLASSES[i]] = float(p.mean()) if p.size > 0 else 0.0
    out["per_class_AP50"] = p50
    return out


@torch.no_grad()
def evaluate(model, processor, val_dl, device):
    model.eval()
    preds = []
    gts = []
    all_image_ids = []
    ann_id = 1
    for batch in val_dl:
        pixel_values = batch["pixel_values"].to(device)
        pixel_mask = batch["pixel_mask"].to(device)
        labels = batch["labels"]
        with torch.autocast(device_type="cuda", dtype=torch.float16):
            outputs = model(pixel_values=pixel_values, pixel_mask=pixel_mask)
        target_sizes = torch.tensor([[labels[bi]["image_size"][0], labels[bi]["image_size"][1]] for bi in range(len(labels))], device=device)
        results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.0)
        for bi, r in enumerate(results):
            lab_gt = labels[bi]
            img_id = int(lab_gt["image_id"])
            all_image_ids.append(img_id)
            scores = r["scores"]
            keep = scores > 0.0
            boxes = r["boxes"][keep]
            scores = scores[keep]
            labels_ids = r["labels"][keep]
            for bx, sc, la in zip(boxes, scores, labels_ids):
                preds.append({"image_id": img_id, "label": int(la.item()), "score": float(sc.item()),
                              "box_xyxy": [float(v) for v in bx]})
            # gt boxes are normalized cxcywh in b['boxes']; convert to pixel xyxy
            lab = labels[bi]
            if lab["boxes"].numel() > 0:
                H, W = int(target_sizes[bi][0]), int(target_sizes[bi][1])
                c = lab["boxes"].float()
                cx, cy = c[:, 0], c[:, 1]
                w2, h2 = c[:, 2], c[:, 3]
                x1 = (cx - w2 / 2) * W
                x2 = (cx + w2 / 2) * W
                y1 = (cy - h2 / 2) * H
                y2 = (cy + h2 / 2) * H
                for j in range(c.shape[0]):
                    x1v, y1v, x2v, y2v = float(x1[j]), float(y1[j]), float(x2[j]), float(y2[j])
                    gts.append({"image_id": img_id, "category_id": int(lab["class_labels"][j].item()),
                                "bbox_xywh": [x1v, y1v, x2v - x1v, y2v - y1v],
                                "area": (x2v - x1v) * (y2v - y1v), "ann_id": ann_id})
                    ann_id += 1
    return to_coco(preds, gts, all_image_ids)


# ---------- main ----------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--epochs", type=int, default=5)
    ap.add_argument("--batch", type=int, default=2)
    ap.add_argument("--size", type=int, default=560)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--workers", type=int, default=2)
    ap.add_argument("--skip_eval", action="store_true")
    ap.add_argument("--no_push", action="store_true")
    args = ap.parse_args()

    tlog = setup_trackio()
    torch.manual_seed(42)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print("device:", device, "gpus:", torch.cuda.device_count(), flush=True)

    processor = DetrImageProcessor.from_pretrained(MODEL_ID)
    processor.do_resize = True
    processor.size = {"shortest_edge": args.size, "longest_edge": 800}
    processor.do_rescale = True
    processor.do_normalize = True
    processor.do_rescale_deprecated = False

    hf_train = load_dataset("biglam/loc_beyond_words", split="train")
    hf_val = load_dataset("biglam/loc_beyond_words", split="validation")
    print(f"train={len(hf_train)} val={len(hf_val)}", flush=True)

    model = DetrForObjectDetection.from_pretrained(
        MODEL_ID, num_labels=7, ignore_mismatched_sizes=True,
        id2label=ID2LABEL, label2id=LABEL2ID,
    ).to(device)
    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print("trainable params:", n_params, flush=True)

    train_ds = DetrDataset(hf_train, processor, "train")
    val_ds = DetrDataset(hf_val, processor, "val")
    train_dl = DataLoader(train_ds, batch_size=args.batch, shuffle=True,
                          num_workers=args.workers, collate_fn=collate_fn, drop_last=True)
    val_dl = DataLoader(val_ds, batch_size=args.batch, shuffle=False,
                        num_workers=args.workers, collate_fn=collate_fn)

    no_decay = ["bias", "LayerNorm.weight", "layer_norm.weight", "embed_positions.weight", "norm.weight"]
    opt = torch.optim.AdamW([
        {"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
         "lr": args.lr, "weight_decay": 1e-4},
        {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
         "lr": args.lr, "weight_decay": 0.0},
    ])
    steps_per_epoch = len(train_dl)
    total_steps = steps_per_epoch * args.epochs
    from transformers import get_linear_schedule_with_warmup
    sched = get_linear_schedule_with_warmup(opt, num_warmup_steps=int(0.1 * total_steps),
                                            num_training_steps=total_steps)

    global_step = 0
    best_map = -1.0
    for epoch in range(args.epochs):
        model.train()
        epoch_loss = 0.0
        nb = 0
        for step, batch in enumerate(train_dl):
            pixel_values = batch["pixel_values"].to(device)
            pixel_mask = batch["pixel_mask"].to(device)
            labels = [{k: (v.to(device) if isinstance(v, torch.Tensor) else v)
                       for k, v in lab.items()} for lab in batch["labels"]]
            with torch.autocast(device_type="cuda", dtype=torch.float16):
                out = model(pixel_values=pixel_values, pixel_mask=pixel_mask, labels=labels)
            loss = sum(v for k, v in out.loss_dict.items() if v is not None)
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            sched.step()
            epoch_loss += loss.item()
            nb += 1
            global_step += 1
            if step % 25 == 0:
                m = {("l_" + k): float(v.item()) for k, v in out.loss_dict.items() if v is not None}
                print(f"[ep{epoch} step{step}/{steps_per_epoch}] loss={loss.item():.3f} { {k: round(v,3) for k,v in m.items()} }", flush=True)
                tlog(step=global_step, loss=loss.item(), **m)
        avg = epoch_loss / max(nb, 1)
        print(f"==== EPOCH {epoch} DONE avg_loss={avg:.4f} ====", flush=True)
        tlog(epoch_loss=avg, epoch=epoch)

        # eval
        if not args.skip_eval:
            print("evaluating...", flush=True)
            metrics = evaluate(model, processor, val_dl, device)
            print("EVAL:", {k: (round(v, 4) if isinstance(v, float) else v) for k, v in metrics.items() if k != "per_class_mAP" and k != "per_class_AP50"}, flush=True)
            print("eval per-class mAP:", {k: round(v, 4) for k, v in metrics["per_class_mAP"].items()}, flush=True)
            print("eval per-class AP50:", {k: round(v, 4) for k, v in metrics["per_class_AP50"].items()}, flush=True)
            tlog(mAP=metrics["mAP"], AP50=metrics["AP50"], AR100=metrics["AR100"], epoch=epoch)
            if metrics["mAP"] > best_map:
                best_map = metrics["mAP"]
                save_dir = "/tmp/best_model"
                model.save_pretrained(save_dir)
                processor.save_pretrained(save_dir)
        else:
            save_dir = "/tmp/best_model"
            model.save_pretrained(save_dir)
            processor.save_pretrained(save_dir)

    print("best mAP:", best_map, flush=True)
    save_dir = "/tmp/final_model"
    model.save_pretrained(save_dir)
    processor.save_pretrained(save_dir)
    print("saved:", save_dir, flush=True)

    if not args.no_push:
        from huggingface_hub import HfApi
        api = HfApi()
        print("pushing model to", REPO_ID, flush=True)
        api.upload_folder(repo_id=REPO_ID, folder_path=save_dir, repo_type="model", commit_message="fine-tuned DETR on loc_beyond_words")
    tlog(best_mAP=best_map)
    try:
        import trackio
        trackio.finish(status=0)
    except Exception:
        pass
    print("DONE", flush=True)


if __name__ == "__main__":
    main()