#!/usr/bin/env python
# /// script
# dependencies = [
#     "torch==2.7.1",
#     "torchvision==0.22.1",
#     "transformers==4.57.6",
#     "timm",
#     "albumentations>=1.4.16,<2.0",
#     "torchmetrics>=1.4",
#     "pycocotools",
#     "datasets>=2.19",
#     "accelerate>=0.34",
#     "safetensors",
#     "huggingface_hub>=0.26",
#     "trackio",
# ]
# ///
"""Fine-tune an object detection model (DETR / RT-DETR family, Apache-2.0 checkpoints)
on biglam/loc_beyond_words and push the result to a hub repo.

Adapted from transformers/examples/pytorch/object-detection/run_object_detection.py
with dataset-schema fixes for biglam/loc_beyond_words (objects["category_id"]),
no-test-split handling and hub push of a proper model card.
"""

import argparse
import json
import logging
import os
import time
from collections.abc import Mapping
from functools import partial
from typing import Any

import albumentations as A
import numpy as np
import torch
from datasets import load_dataset
from torchmetrics.detection.mean_ap import MeanAveragePrecision

from transformers import (
    AutoConfig,
    AutoImageProcessor,
    AutoModelForObjectDetection,
    Trainer,
    TrainingArguments,
)
from transformers.image_processing_utils import BatchFeature
from transformers.image_transforms import center_to_corners_format
from transformers.trainer import EvalPrediction

logger = logging.getLogger(__name__)
logging.basicConfig(format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
                    datefmt="%m/%d/%Y %H:%M:%S", handlers=[logging.StreamHandler()], level=logging.INFO)


def format_image_annotations_as_coco(image_id, categories, areas, bboxes):
    annotations = []
    for category, area, bbox in zip(categories, areas, bboxes):
        annotations.append({
            "image_id": image_id,
            "category_id": category,
            "iscrowd": 0,
            "area": area,
            "bbox": list(bbox),
        })
    return {"image_id": image_id, "annotations": annotations}


class ModelOutput:
    def __init__(self, logits, pred_boxes):
        self.logits = logits
        self.pred_boxes = pred_boxes


def convert_bbox_yolo_to_pascal(boxes, image_size):
    boxes = center_to_corners_format(boxes)
    height, width = image_size
    boxes = boxes * torch.tensor([[width, height, width, height]])
    return boxes


def augment_and_transform_batch(examples, transform, image_processor, return_pixel_mask=False):
    # biglam/loc_beyond_words stores "objects" as a List of dicts per row (not a nested
    # Sequence dict), so re-aggregate bboxes / categories / areas per row here.
    images, annotations = [], []
    for image_id, image, objects in zip(examples["image_id"], examples["image"], examples["objects"]):
        image = np.array(image.convert("RGB"))
        bboxes = [o["bbox"] for o in objects]
        cats = [o["category_id"] for o in objects]
        areas = [o["area"] for o in objects]
        output = transform(image=image, bboxes=bboxes, category=cats)
        images.append(output["image"])
        formatted = format_image_annotations_as_coco(
            image_id, output["category"], areas, output["bboxes"]
        )
        annotations.append(formatted)
    result = image_processor(images=images, annotations=annotations, return_tensors="pt")
    if not return_pixel_mask:
        result.pop("pixel_mask", None)
    return result


def collate_fn(batch):
    data = {}
    data["pixel_values"] = torch.stack([x["pixel_values"] for x in batch])
    data["labels"] = [x["labels"] for x in batch]
    if "pixel_mask" in batch[0]:
        data["pixel_mask"] = torch.stack([x["pixel_mask"] for x in batch])
    return data


@torch.no_grad()
def compute_metrics(evaluation_results, image_processor, threshold=0.0, id2label=None, eval_gts=None):
    """COCO-style metrics from a Trainer eval.

    eval_gts: list of per-sample dicts precomputed in the exact order of the
    (unshuffled) validation dataset: {"orig_size": [H, W],
    "boxes_xyxy": (n,4) absolute pixel boxes, "labels": (n,) int labels}.
    """
    predictions = evaluation_results.predictions

    # Locate logits / pred_boxes. RT-DETR's forward returns a tuple of many
    # arrays when labels are passed (losses + logits + boxes); DETR returns
    # (loss, logits, boxes). Identify by shape.
    logits = boxes = None
    if isinstance(predictions, Mapping) and "logits" in predictions:
        logits = torch.as_tensor(predictions["logits"])
        boxes = torch.as_tensor(predictions["pred_boxes"])
    else:
        for arr in predictions:
            if hasattr(arr, "ndim") and arr.ndim == 3:
                t = torch.as_tensor(arr)
                if arr.shape[-1] == 4 and boxes is None:
                    boxes = t
                elif arr.shape[-1] > 4 and logits is None:
                    logits = t
        if logits is None or boxes is None:
            raise RuntimeError("could not locate logits/pred_boxes in eval predictions")

    n = logits.shape[0]
    if eval_gts is None:
        raise RuntimeError("compute_metrics requires precomputed eval_gts")
    gts = eval_gts[:n]  # guard against trainer dropping tail samples
    target_sizes = torch.tensor([g["orig_size"] for g in gts])

    output = ModelOutput(logits=logits, pred_boxes=boxes)
    post_processed_predictions = image_processor.post_process_object_detection(
        output, threshold=threshold, target_sizes=target_sizes
    )
    post_processed_targets = [
        {"boxes": torch.tensor(g["boxes_xyxy"]), "labels": torch.tensor(g["labels"])} for g in gts
    ]

    metric = MeanAveragePrecision(box_format="xyxy", class_metrics=True)
    metric.update(post_processed_predictions, post_processed_targets)
    metrics = metric.compute()

    classes = metrics.pop("classes")
    map_per_class = metrics.pop("map_per_class")
    mar_100_per_class = metrics.pop("mar_100_per_class")
    for class_id, class_map, class_mar in zip(classes, map_per_class, mar_100_per_class):
        class_name = id2label[class_id.item()] if id2label is not None else str(class_id.item())
        metrics[f"map_{class_name}"] = class_map
        metrics[f"mar_100_{class_name}"] = class_mar
    return {k: round(v.item(), 4) for k, v in metrics.items()}


def build_parser():
    p = argparse.ArgumentParser()
    p.add_argument("--dataset-name", default="biglam/loc_beyond_words")
    p.add_argument("--model-name-or-path", default="PekingU/rtdetr_r18vd")
    p.add_argument("--image-square-size", type=int, default=640)
    p.add_argument("--epochs", type=int, default=30)
    p.add_argument("--batch-size", type=int, default=8)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--weight-decay", type=float, default=1e-4)
    p.add_argument("--warmup-steps", type=int, default=200)
    p.add_argument("--grad-accum", type=int, default=1)
    p.add_argument("--eval-steps", type=int, default=712)
    p.add_argument("--save-steps", type=int, default=712)
    p.add_argument("--save-total-limit", type=int, default=3)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--max-train-samples", type=int, default=None)
    p.add_argument("--max-eval-samples", type=int, default=None)
    p.add_argument("--num-workers", type=int, default=4)
    p.add_argument("--output-dir", default="/root/output")
    p.add_argument("--hub-repo", default="harness-race/prime-r2")
    p.add_argument("--push", action="store_true")
    p.add_argument("--profile-batches", type=int, default=0,
                   help="if >0: time N training batches, print throughput, exit without training")
    p.add_argument("--disable-augmentations", action="store_true")
    p.add_argument("--eval-max-batches", type=int, default=None,
                   help="cap eval batches for profiling")
    p.add_argument("--eval-accumulation-steps", type=int, default=4,
                   help="offload eval predictions to CPU every N batches (avoids GPU OOM)")
    p.add_argument("--mini-train", action="store_true",
                   help="after profiling, also run a short real training+eval phase")
    return p


def main():
    args = build_parser().parse_args()

    # ---- optional trackio logging (wrapped, must never break training) ----
    track = None
    try:
        import trackio as _t
        _t.init(project="prime-r2-loc-beyond-words", name=os.path.basename(args.model_name_or_path))
        track = _t
        logger.info("trackio logging enabled")
    except Exception as e:
        logger.info(f"trackio disabled: {e}")

    dataset = load_dataset(args.dataset_name)
    if "validation" not in dataset:
        split = dataset["train"].train_test_split(0.15, seed=args.seed)
        dataset["train"] = split["train"]
        dataset["validation"] = split["test"]
    if args.max_train_samples:
        dataset["train"] = dataset["train"].select(range(args.max_train_samples))
    if args.max_eval_samples:
        dataset["validation"] = dataset["validation"].select(range(args.max_eval_samples))

    feats = dataset["train"].features["objects"]
    if isinstance(feats, dict):
        categories = feats["category_id"].feature.names
    else:
        categories = feats.feature["category_id"].names
    id2label = dict(enumerate(categories))
    label2id = {v: k for k, v in id2label.items()}
    logger.info(f"classes ({len(categories)}): {id2label}")

    config = AutoConfig.from_pretrained(args.model_name_or_path, label2id=label2id, id2label=id2label)
    model = AutoModelForObjectDetection.from_pretrained(
        args.model_name_or_path, config=config, ignore_mismatched_sizes=True
    )
    image_processor = AutoImageProcessor.from_pretrained(
        args.model_name_or_path,
        do_resize=True,
        size={"max_height": args.image_square_size, "max_width": args.image_square_size},
        do_pad=True,
        pad_size={"height": args.image_square_size, "width": args.image_square_size},
        use_fast=False,
    )

    max_size = args.image_square_size
    bbox_params = A.BboxParams(format="coco", label_fields=["category"], clip=True, min_area=25)
    if args.disable_augmentations:
        train_transform = A.Compose([A.NoOp()], bbox_params=bbox_params)
    else:
        train_transform = A.Compose(
            [
                A.Compose(
                    [A.SmallestMaxSize(max_size=max_size, p=1.0),
                     A.RandomSizedBBoxSafeCrop(height=max_size, width=max_size, p=1.0)],
                    p=0.2,
                ),
                A.OneOf(
                    [A.Blur(blur_limit=7, p=0.5), A.MotionBlur(blur_limit=7, p=0.5)],
                    p=0.1,
                ),
                A.Perspective(p=0.1),
                A.HorizontalFlip(p=0.5),
                A.RandomBrightnessContrast(p=0.5),
            ],
            bbox_params=bbox_params,
        )
    validation_transform = A.Compose([A.NoOp()], bbox_params=bbox_params)

    train_transform_batch = partial(augment_and_transform_batch, transform=train_transform,
                                    image_processor=image_processor)
    validation_transform_batch = partial(augment_and_transform_batch, transform=validation_transform,
                                         image_processor=image_processor)
    dataset["train"] = dataset["train"].with_transform(train_transform_batch)
    dataset["validation"] = dataset["validation"].with_transform(validation_transform_batch)

    # GPU?
    device = "cuda" if torch.cuda.is_available() else "cpu"
    logger.info(f"device={device} ({torch.cuda.get_device_name(0) if device=='cuda' else 'n/a'})")
    logger.info(f"model params: {sum(p.numel() for p in model.parameters())/1e6:.1f}M")

    # ---- quick profiling mode: time N training batches ----
    if args.profile_batches > 0:
        model = model.to(device)
        model.train()
        dl = torch.utils.data.DataLoader(
            dataset["train"], batch_size=args.batch_size, collate_fn=collate_fn,
            num_workers=args.num_workers, shuffle=True
        )
        def _to_dev(v):
            if isinstance(v, torch.Tensor):
                return v.to(device)
            if isinstance(v, (list, tuple)):
                return type(v)(_to_dev(x) for x in v)
            if isinstance(v, Mapping):
                return {k: _to_dev(x) for k, x in v.items()}
            return v
        batch = next(iter(dl))
        batch = {k: _to_dev(v) for k, v in batch.items()}
        # warmup
        for _ in range(3):
            loss = model(**batch).loss
            loss.backward()
            model.zero_grad()
        torch.cuda.synchronize() if device == "cuda" else None
        t0 = time.time()
        for _ in range(args.profile_batches):
            loss = model(**batch).loss
            loss.backward()
            model.zero_grad()
        torch.cuda.synchronize() if device == "cuda" else None
        dt = (time.time() - t0) / args.profile_batches
        n = len(dataset["train"])
        per_epoch = n / args.batch_size
        info = {
            "seconds_per_batch": round(dt, 3),
            "batches_per_sec": round(1.0 / dt, 3),
            "train_examples": n,
            "steps_per_epoch": per_epoch,
            "est_seconds_per_epoch": round(dt * per_epoch, 1),
            "profile_batches": args.profile_batches,
            "batch_size": args.batch_size,
        }
        print("PROFILE_JSON " + json.dumps(info))
        logger.info("PROFILE_JSON " + json.dumps(info))
        if not args.mini_train:
            return

    # ---- training ----
    train_args = TrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        weight_decay=args.weight_decay,
        warmup_steps=args.warmup_steps,
        lr_scheduler_type="cosine",
        fp16=(device == "cuda"),
        bf16=False,
        dataloader_num_workers=args.num_workers,
        dataloader_pin_memory=True,
        remove_unused_columns=False,
        eval_strategy="steps",
        eval_steps=args.eval_steps,
        eval_accumulation_steps=args.eval_accumulation_steps,
        logging_steps=20,
        save_strategy="steps",
        save_steps=args.save_steps,
        save_total_limit=args.save_total_limit,
        load_best_model_at_end=True,
        metric_for_best_model="map",
        greater_is_better=True,
        seed=args.seed,
        report_to=[],
        run_name="prime-r2-loc-beyond-words",
        ddp_find_unused_parameters=None,
        group_by_length=False,
    )

    # Precompute eval ground truth once (deterministic, NoOp transform) so metrics don't
    # depend on the trainer's label flattening. Order = validation dataset order = eval order.
    eval_gts = []
    for ex in dataset["validation"]:
        labels = ex["labels"]
        boxes = torch.tensor(labels["boxes"])  # (n,4) normalized cxcywh
        h, w = labels["orig_size"][0].item(), labels["orig_size"][1].item()
        boxes = convert_bbox_yolo_to_pascal(boxes, (h, w))
        eval_gts.append({
            "orig_size": [int(h), int(w)],
            "boxes_xyxy": boxes.numpy().tolist(),
            "labels": np.asarray(labels["class_labels"], dtype=np.int64),
        })
    logger.info(f"precomputed eval ground truth for {len(eval_gts)} images")

    eval_compute_metrics_fn = partial(compute_metrics, image_processor=image_processor,
                                      id2label=id2label, threshold=0.0, eval_gts=eval_gts)

    trainer = Trainer(
        model=model,
        args=train_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["validation"],
        processing_class=image_processor,
        data_collator=collate_fn,
        compute_metrics=eval_compute_metrics_fn,
    )

    if args.eval_max_batches:
        # quick eval sanity: evaluate only first N batches before training
        small = dataset["validation"].select(range(args.batch_size * args.eval_max_batches))
        logger.info("pre-training sanity eval on %d images", len(small))
        metrics0 = trainer.evaluate(eval_dataset=small, metric_key_prefix="init")
        logger.info("INIT_EVAL %s", json.dumps(metrics0, default=str))

    train_result = trainer.train()
    trainer.save_model(os.path.join(args.output_dir, "final_model"))
    trainer.log_metrics("train", train_result.metrics)
    trainer.save_metrics("train", train_result.metrics)
    trainer.save_state()

    metrics = trainer.evaluate(metric_key_prefix="test")
    trainer.log_metrics("test", metrics)
    trainer.save_metrics("test", metrics)

    results = {
        "train": train_result.metrics,
        "test": metrics,
        "model": args.model_name_or_path,
        "dataset": args.dataset_name,
        "image_square_size": args.image_square_size,
        "epochs": args.epochs,
        "batch_size": args.batch_size,
        "lr": args.lr,
        "id2label": id2label,
    }
    with open(os.path.join(args.output_dir, "metrics.json"), "w") as f:
        json.dump(results, f, indent=2)
    logger.info("FINAL_METRICS " + json.dumps(metrics))
    logger.info("ALL_RESULTS_JSON " + json.dumps(results, default=str))
    if track is not None:
        try:
            track.log({k: v for k, v in metrics.items() if isinstance(v, (int, float))})
            track.finish()
        except Exception as e:
            logger.info(f"trackio log failed: {e}")

    # ---- push to hub ----
    if args.push:
        from huggingface_hub import HfApi
        api = HfApi(token=os.environ.get("HF_TOKEN"))
        repo_id = args.hub_repo
        api.create_repo(repo_id, repo_type="model", exist_ok=True)
        final_dir = os.path.join(args.output_dir, "final_model")
        # write README model card
        readme = build_model_card(results, id2label)
        with open(os.path.join(final_dir, "README.md"), "w") as f:
            f.write(readme)
        api.upload_folder(
            folder_path=final_dir,
            repo_id=repo_id,
            repo_type="model",
            commit_message=f"Fine-tune {args.model_name_or_path} on {args.dataset_name} (7 classes)",
        )
        logger.info(f"Pushed model to {repo_id}")
        print(f"PUSHED {repo_id}")


def build_model_card(results, id2label):
    metric_rows = []
    keys = ["map", "map_50", "map_75", "mar_1", "mar_10", "mar_100"]
    m = results["test"]
    for k in keys:
        if k in m:
            metric_rows.append(f"| {k} | {m[k]} |")
    per_class = ""
    for cid, name in id2label.items():
        mk = f"map_{name}"
        if f"map_{name}" in m:
            per_class += f"| {name} | {m[f'map_{name}']} |\n"
    base = results["model"]
    ds = results["dataset"]
    model_card = f"""---
license: apache-2.0
base_model: {base}
tags:
- object-detection
- vision
- transformers
- pytorch
- document-layout-analysis
- newspapers
datasets:
- {ds}
metrics:
- {', '.join([k for k in ['map','map_50','map_75','mar_100'] if k in results['test']])}
pipeline_tag: object-detection
---

# Prime R2 — Object Detection on LOC Beyond Words

Fine-tuned object detection model for the [**Beyond Words**](https://huggingface.co/datasets/{ds}) newspaper page
layout dataset (Library of Congress / biglam). Detects 7 element types in digitized newspaper pages:

{", ".join([f"**{n}**" for n in id2label.values()])}

## Model

- **Base model:** [`{base}`](https://huggingface.co/{base}) — license: **Apache-2.0** (open, shareable)
- **Architecture:** Transformers `AutoModelForObjectDetection` (DETR/RT-DETR family)
- **Input:** grayscale newspaper page images converted to RGB, resized/padded to {results.get('image_square_size', '?')}×{results.get('image_square_size', '?')}
- **Bounding boxes:** COCO format (x, y, width, height)

## Training

- **Dataset:** [{ds}](https://huggingface.co/datasets/{ds}) (CC0-1.0) — 2,846 train / 712 validation images
- **Epochs:** {results.get('epochs')}, **batch size:** {results.get('batch_size')}, **learning rate:** {results.get('lr')}
- **Optimizer:** AdamW, cosine schedule, warmup; FP16 mixed precision (GPU job)
- **Augmentations:** random sized bbox-safe crop, blur, perspective, horizontal flip, brightness/contrast

## Validation results (COCO metrics, torchmetrics, threshold 0.0)

| Metric | Value |
|---|---|
{''.join(metric_rows)}

### Per-class mAP

| Class | mAP |
|---|---|
{per_class}

## How to use

```python
from transformers import AutoModelForObjectDetection, AutoImageProcessor
import torch

model = AutoModelForObjectDetection.from_pretrained("harness-race/prime-r2")
processor = AutoImageProcessor.from_pretrained("harness-race/prime-r2")

image = <PIL.Image in RGB>
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)
results = processor.post_process_object_detection(
    outputs, threshold=0.5, target_sizes=torch.tensor([image.size[::-1]])
)
for r in results[0]:
    print(r["label"], r["score"], r["box"])
```

## License
Apache-2.0 (model weights). Dataset is CC0-1.0.
"""
    return model_card


if __name__ == "__main__":
    main()