# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch>=2.1",
#     "transformers>=4.40",
#     "accelerate>=0.27",
#     "datasets>=2.18",
#     "pycocotools",
#     "huggingface_hub>=0.23",
#     "Pillow",
#     "numpy",
#     "timm",
#     "scipy",
#     "requests",
# ]
# ///
import os, json, time, argparse, io
import numpy as np
import torch
from PIL import Image

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--data_dir", default="/data")
    p.add_argument("--base_model", default="facebook/detr-resnet-50")
    p.add_argument("--hub_id", default="harness-race/opencode-r3")
    p.add_argument("--epochs", type=int, default=25)
    p.add_argument("--batch_size", type=int, default=4)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--patience", type=int, default=4)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--exp_name", default="opencode-r3-detr")
    p.add_argument("--max_steps", type=int, default=0)
    return p.parse_args()

ARGS = parse_args()

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()}
N = len(ID2LABEL)

import datasets as hfds
from torch.utils.data import Dataset, DataLoader
from transformers import DetrImageProcessor, DetrForObjectDetection
from tqdm import tqdm
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval

def log(*m):
    print("[%s]" % time.strftime("%H:%M:%S"), *m, flush=True)

torch.manual_seed(ARGS.seed)
np.random.seed(ARGS.seed)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
log("torch", torch.__version__, "device", device, "cuda_avail", torch.cuda.is_available())

data_dir = ARGS.data_dir
if not os.path.isdir(data_dir):
    log("mount %s not found; downloading dataset" % data_dir)
    data_dir = ARGS.data_dir = "biglam/loc_beyond_words"

try:
    ds = hfds.load_dataset(data_dir, split="train")
    ds_val = hfds.load_dataset(data_dir, split="validation")
except Exception as e:
    log("local load failed (%s); falling back to hub dataset" % e)
    data_dir = "biglam/loc_beyond_words"
    ds = hfds.load_dataset(data_dir, split="train")
    ds_val = hfds.load_dataset(data_dir, split="validation")

def get_pil(ex):
    img = ex["image"]
    if isinstance(img, dict):
        img = Image.open(io.BytesIO(img["bytes"])).convert("RGB")
    else:
        img = img.convert("RGB")
    return img

def fields(ex):
    objs = ex["objects"]
    bboxes=[]; cats=[]; areas=[]; ids=[]
    if objs is None: objs=[]
    for o in objs:
        b = o["bbox"]
        if isinstance(b, (str, tuple)):
            b = list(float(x) for x in b)
        bboxes.append([float(x) for x in b])
        c = o["category_id"]
        if isinstance(c, str):
            cats.append(LABEL2ID.get(c, 0))
        else:
            try: cats.append(int(c)&0x7fffffff)
            except Exception: cats.append(0)
        areas.append(float(o.get("area", b[2]*b[3])))
        ids.append(int(o.get("id", 0)))
    return {"bbox":bboxes,"category_id":cats,"area":areas,"id":ids}

# Build torch Dataset returning (pil_img, anns_for_processor, size_wh, row_id)
class DetrDS(Dataset):
    def __init__(self, examples, with_size=True):
        self.examples = examples
        self.sizes = []
        for ex in examples:
            try:
                w,h = ex["width"], ex["height"]
            except Exception:
                w,h = get_pil(ex).size
            self.sizes.append((w,h))
    def __len__(self):
        return len(self.examples)
    def __getitem__(self, i):
        ex = self.examples[i]
        img = get_pil(ex)
        o = fields(ex)
        bboxes = o.get("bbox") or []
        cats = o.get("category_id")
        if not isinstance(cats, list):
            cats = [0]*len(bboxes)
        anns = [{"bbox":[float(x) for x in bboxes[j]], "category_id":int(cats[j])&0x7fffffff,
                 "area": float(bboxes[j][2]*bboxes[j][3])} for j in range(len(bboxes))]
        return img, anns, self.sizes[i], i

processor = DetrImageProcessor.from_pretrained(ARGS.base_model,
    do_resize=True, size={"shortest_edge":800, "longest_edge":1333})

train_ds = DetrDS(list(ds))
val_ds = DetrDS(list(ds_val))
log("train", len(train_ds), "val", len(val_ds))

def collate_fn(batch):
    # batch: list of (img, anns, size, idx)
    images = [b[0] for b in batch]
    anns = [b[1] for b in batch]
    sizes = [b[2] for b in batch]
    idx = [b[3] for b in batch]
    return images, anns, sizes, idx

train_dl = DataLoader(train_ds, batch_size=ARGS.batch_size, shuffle=True, num_workers=2, collate_fn=collate_fn)
val_dl = DataLoader(val_ds, batch_size=ARGS.batch_size, shuffle=False, num_workers=2, collate_fn=collate_fn)

model = DetrForObjectDetection.from_pretrained(ARGS.base_model, num_labels=N,
    ignore_mismatched_sizes=True, id2label=ID2LABEL, label2id=LABEL2ID)
model.to(device)

# distinct backbone (frozen-ish) vs head lr
bb_params   = [p for n,p in model.model.backbone.named_parameters() if "layers" in n]
head_params = [p for n,p in model.named_parameters() if not n.startswith("model.backbone") or "layers" not in n]
optim = torch.optim.AdamW([
    {"params": bb_params, "lr": ARGS.lr/10},
    {"params": head_params, "lr": ARGS.lr},
], lr=ARGS.lr, weight_decay=1e-4)
scaler = torch.amp.GradScaler("cuda", init_scale=2.0**12) if device.type=="cuda" else None

def build_coco_gt(examples):
    gt = {"images":[],"annotations":[],"categories":[{"id":k,"name":v} for k,v in ID2LABEL.items()]}
    ann_id = 1
    for i, ex in enumerate(examples):
        o = fields(ex)
        w,h = val_ds.sizes[i]
        gt["images"].append({"id":i,"width":w,"height":h})
        bboxes = o.get("bbox") or []
        cats = o.get("category_id")
        if not isinstance(cats, list):
            cats = [0]*len(bboxes)
        for j in range(len(bboxes)):
            b = [float(x) for x in bboxes[j]]
            gt["annotations"].append({"id":ann_id,"image_id":i,"category_id":int(cats[j])&0x7fffffff,
                                      "bbox":b,"area":b[2]*b[3],"iscrowd":0}); ann_id+=1
    g = COCO(); g.dataset = gt; g.createIndex()
    return g

coco_gt = build_coco_gt(list(ds_val))

def evaluate():
    model.eval()
    dets = []
    ann_id = 1
    with torch.no_grad():
        for images, anns, sizes, idx in tqdm(val_dl, desc="eval"):
            enc = processor(images=images, return_tensors="pt")
            pv = enc["pixel_values"].to(device); pm = enc["pixel_mask"].to(device)
            with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=(device.type=="cuda")):
                out = model(pixel_values=pv, pixel_mask=pm)
            target_sizes = torch.tensor([val_ds.sizes[ii][::-1] for ii in idx.astype(int)] if hasattr(idx,'astype') else [[val_ds.sizes[ii][1], val_ds.sizes[ii][0]] for ii in idx], device=device)
            # idx is list of ints already
            sz = [[val_ds.sizes[ii][1], val_ds.sizes[ii][0]] for ii in idx]
            results = processor.post_process_object_detection(out, threshold=0.0, target_sizes=sz)
            for ii, res in zip(idx, results):
                boxes = res["boxes"].cpu().numpy()   # xmin ymin xmax ymax
                scores = res["scores"].cpu().numpy()
                labels = res["labels"].cpu().numpy()
                for b,s,l in zip(boxes,scores,labels):
                    if s <= 0.0: continue
                    dets.append({"id":ann_id,"image_id":ii,"category_id":int(l),"bbox":[float(b[0]),float(b[1]),float(b[2]-b[0]),float(b[3]-b[1])],"score":float(s)})
                    ann_id += 1
    model.train()
    if len(dets)==0:
        return None
    try:
        preds = coco_gt.loadRes(dets)
        evaluator = COCOeval(coco_gt, preds, iouType="bbox")
        # set areas to all, max dets high
        evaluator.params.maxDets = [100, 300, 1000]
        evaluator.evaluate(); evaluator.accumulate(); evaluator.summarize()
        ap = evaluator.stats  # [0]=AP@.5:.95 ... [5]=AP50 ... [6]=AR100
        return {"AP": float(ap[0]), "AP50": float(ap[1]), "AP75": float(ap[2]),
                "AP_s": float(ap[3]), "AP_m": float(ap[4]), "AP_l": float(ap[5]),
                "AR_max100": float(ap[6]), "AR_max1000": float(ap[8])}
    except Exception as e:
        log("COCOeval failed:", e)
        return None

# ---------- training ----------
best = {"AP": -1.0}
no_improve = 0
global_step = 0
run_outputs = {}
t_start = time.time()

log("starting training; epochs", ARGS.epochs, "batch", ARGS.batch_size)
for epoch in range(1, ARGS.epochs+1):
    model.train()
    ep_loss = 0.0; ep_losssum = {}; nb = 0
    for images, anns, sizes, idx in tqdm(train_dl, desc="epoch %d"%epoch):
        anns_pp = [{"image_id": int(idx[k]), "annotations": anns[k]} for k in range(len(images))]
        enc = processor(images=images, annotations=anns_pp, return_tensors="pt")
        pv = enc["pixel_values"].to(device); pm = enc["pixel_mask"].to(device)
        labels = [{k: v.to(device) for k,v in l.items()} for l in enc.get("labels")]
        optim.zero_grad()
        try:
            with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=(device.type=="cuda")):
                out = model(pixel_values=pv, pixel_mask=pm, labels=labels)
                loss = out.loss
        except Exception as e:
            log("skipping batch that raised:", e)
            continue
        if not torch.isfinite(loss):
            log("skipping non-finite loss step")
            continue
        scaler.scale(loss).backward()
        scaler.step(optim)
        scaler.update()
        ep_loss += float(loss.detach().float())
        for kk, vv in out.loss_dict.items():
            ep_losssum[kk] = ep_losssum.get(kk, 0.0) + float(vv.detach().float())
        nb += 1
        global_step += 1
        if ARGS.max_steps and global_step >= ARGS.max_steps:
            break
    log("epoch %d loss %.4f %s" % (epoch, ep_loss/nb,
        " ".join("%s %.4f"%(k, v/nb) for k,v in ep_losssum.items())))
    m = evaluate()
    log("VAL", json.dumps(m))
    if m and m["AP"] > best["AP"]:
        best = m
        best.update({"epoch": epoch, "global_step": global_step})
        no_improve = 0
        torch.save({"state_dict": model.state_dict()}, os.path.join(os.getcwd(), "best_model.pt"))
        log("saved new best AP=%.4f" % m["AP"])
    else:
        no_improve += 1
        if no_improve >= ARGS.patience:
            log("early stop after epoch", epoch)
            break
    if ARGS.max_steps and global_step >= ARGS.max_steps:
        break
    # decay lr slowly
    # (optional) torch.optim.lr_scheduler not added; keep fixed

run_outputs["best"] = best
run_outputs["epochs_run"] = epoch
run_outputs["elapsed_sec"] = round(time.time()-t_start, 1)
log("best val result:", best)

# ---------- reload best & push ----------
if os.path.exists("best_model.pt"):
    sd = torch.load("best_model.pt", map_location="cpu")
    model.load_state_dict(sd["state_dict"])
else:
    best = evaluate() or {}
    run_outputs["best"] = best

    model.push_to_hub(ARGS.hub_id)
    processor.push_to_hub(ARGS.hub_id)
log("pushed weights + preprocessor")

# build model card
card = {
    "library_name": "transformers",
    "pipeline_tag": "object-detection",
    "license": "apache-2.0",
    "tags": ["object-detection","detr","computer-vision","document-layout-analysis","pytorch"],
    "base_model": ARGS.base_model,
    "model-index": [{
        "name": ARGS.exp_name,
        "results": [{
            "task": {"type":"object-detection"},
            "dataset": {"type":"biglam/loc_beyond_words","name":"Beyond Words (Testing)","config":"default"},
            "metrics": [
                {"type":"Average Precision","name":"mAP @[IoU=0.50:0.95]","value": round(run_outputs["best"].get("AP",0.0),4)},
                {"type":"Average Precision","name":"mAP @[IoU=0.50]","value": round(run_outputs["best"].get("AP50",0.0),4)},
            ]
        }]
    }]
}


def build_readme(res):
    m = res if res else {}
    class Row:
        pass
    def f(k, d=.0):
        return "%.4f" % m.get(k, d)
    front = json.dumps(card, indent=2)
    return f"""---
{front}
---

# opencode-r3: Beyond Words Object Detection (DETR-ResNet-50)

Model for detecting visual content regions in WWI-era historical newspaper pages
from the US Library of Congress **Beyond Words / Chronicling America** collection.
The base model `{ARGS.base_model}` (Apache-2.0) was fine-tuned on {len(train_ds)} train images
and evaluated on {len(val_ds)} validation images over **7 classes**:
Photograph, Illustration, Map, Comics/Cartoon, Editorial Cartoon, Headline, Advertisement.

## Model Details

| | |
|---|---|
| **Base model** | {ARGS.base_model} (DETR, ResNet-50 backbone) |
| **License** | Apache-2.0 (open, shareable) |
| **Architecture** | `DetrForObjectDetection` |
| **Task** | Object detection / document layout analysis |
| **Dataset** | [biglam/loc_beyond_words](https://huggingface.co/datasets/biglam/loc_beyond_words) (CC0-1.0) |
| **Splits** | {len(train_ds)} train / {len(val_ds)} validation |

## Training Procedure

- **Optimizer:** AdamW (head LR {ARGS.lr}, backbone LR {ARGS.lr/10:g})
- **Epochs:** {run_outputs.get('epochs_run','?')} (early-stopped, patience {ARGS.patience})
- **Batch size:** {ARGS.batch_size}
- **Image size:** longest edge 1333 / shortest edge 800
- **Hardware:** Hugging Face Jobs GPU (T4)
- **Selection:** best checkpoint by validation mAP@[0.50:0.95]

## Evaluation Results (validation)

| Metric | Value |
|---|---|
| **mAP @[IoU=0.50:0.95]** | **{f('AP')}** |
| mAP @[IoU=0.50] | {f('AP50')} |
| mAP @[IoU=0.75] | {f('AP75')} |
| AP small | {f('AP_s')} |
| AP medium | {f('AP_m')} |
| AP large | {f('AP_l')} |
| AR (maxDets=100) | {f('AR_max100')} |

## Usage

```python
import torch
from transformers import DetrImageProcessor, DetrForObjectDetection
from PIL import Image

model = DetrForObjectDetection.from_pretrained("harness-race/opencode-r3")
processor = DetrImageProcessor.from_pretrained("harness-race/opencode-r3")
image = Image.open("page.jpg").convert("RGB")
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)
results = processor.post_process_object_detection(
    outputs, target_sizes=[image.size[::-1]], threshold=0.5)[0]
for score, label_id, box in zip(results["scores"], results["labels"], results["boxes"]):
    if score > 0.5:
        print(model.config.id2label[label_id.item()], round(score.item(),3), box.tolist())
```

## Licensing

Base model `{ARGS.base_model}` is **Apache-2.0** (permissive open license), so this
fine-tuned model may be freely shared and reused. The training dataset is public
domain (**CC0-1.0**).

## Known Limitations

- Trained on a single era/language; pre-1875 layouts may underperform.
- Skewed class distribution (headlines/ads dominate) can depress rare-class AP.
"""

readme = build_readme(run_outputs["best"])
with open("README.md","w") as f:
    f.write(readme)

# push README + a results.json
from huggingface_hub import HfApi
api = HfApi()
with open("validation_results.json", "w") as f:
    json.dump(run_outputs["best"], f, indent=2)
with open("README.md", "w") as f:
    f.write(readme)
api.upload_file(path_or_fileobj="validation_results.json", path_in_repo="validation_results.json", repo_id=ARGS.hub_id, repo_type="model")
api.upload_file(path_or_fileobj="README.md", path_in_repo="README.md", repo_id=ARGS.hub_id, repo_type="model")
# model card metadata (config fields like license) already in config.json via push_to_hub

log("=== FINAL VALIDATION RESULTS ===")
print(json.dumps(run_outputs["best"], indent=2))
log("done")