import argparse
import json
from pathlib import Path

import torch
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm
from transformers import RTDetrForObjectDetection, RTDetrImageProcessor

BASE_MODEL = "PekingU/rtdetr_r50vd"

def classes(path):
    return [x.strip() for x in Path(path).read_text().splitlines() if x.strip()]

class COCODataset(Dataset):
    def __init__(self, image_dir, ann_file, processor):
        self.image_dir = Path(image_dir)
        self.processor = processor
        coco = json.loads(Path(ann_file).read_text())
        cats = sorted(coco["categories"], key=lambda x: x["id"])
        self.cat_map = {c["id"]: i for i, c in enumerate(cats)}
        anns = {}
        for a in coco["annotations"]:
            anns.setdefault(a["image_id"], []).append(a)
        self.records = []
        for im in coco["images"]:
            self.records.append({
                "id": im["id"],
                "file_name": im["file_name"],
                "annotations": anns.get(im["id"], []),
            })

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

    def __getitem__(self, i):
        r = self.records[i]
        image = Image.open(self.image_dir / r["file_name"]).convert("RGB")
        anns = []
        for a in r["annotations"]:
            x, y, w, h = a["bbox"]
            anns.append({
                "id": a["id"],
                "image_id": r["id"],
                "category_id": self.cat_map[a["category_id"]],
                "bbox": [x, y, w, h],
                "area": float(a["area"]),
                "iscrowd": 0,
            })
        out = self.processor(
            images=image,
            annotations={
                "image_id": r["id"],
                "annotations": anns,
            },
            return_tensors="pt",
        )
        out["pixel_values"] = out["pixel_values"].squeeze(0)
        if "pixel_mask" in out:
            out["pixel_mask"] = out["pixel_mask"].squeeze(0)
        out["labels"] = out["labels"][0]
        return out

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

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--train-dir", required=True)
    p.add_argument("--val-dir", required=True)
    p.add_argument("--classes", required=True)
    p.add_argument("--output-dir", required=True)
    p.add_argument("--epochs", type=int, default=30)
    p.add_argument("--batch-size", type=int, default=2)
    args = p.parse_args()

    names = classes(args.classes)
    id2label = {i: n for i, n in enumerate(names)}
    label2id = {n: i for i, n in enumerate(names)}

    processor = RTDetrImageProcessor.from_pretrained(BASE_MODEL)

    train = COCODataset(
        Path(args.train_dir)/"images",
        Path(args.train_dir)/"annotations.json",
        processor,
    )
    val = COCODataset(
        Path(args.val_dir)/"images",
        Path(args.val_dir)/"annotations.json",
        processor,
    )

    model = RTDetrForObjectDetection.from_pretrained(
        BASE_MODEL,
        num_labels=len(names),
        id2label=id2label,
        label2id=label2id,
        ignore_mismatched_sizes=True,
    )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    train_loader = DataLoader(
        train, batch_size=args.batch_size, shuffle=True,
        num_workers=0, collate_fn=collate
    )

    val_loader = DataLoader(
        val, batch_size=args.batch_size, shuffle=False,
        num_workers=0, collate_fn=collate
    )

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

    outdir = Path(args.output_dir)
    outdir.mkdir(parents=True, exist_ok=True)
    best = float("inf")

    for epoch in range(args.epochs):
        model.train()
        running = 0.0

        for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}"):
            batch = {
                k: (v.to(device) if torch.is_tensor(v) else v)
                for k, v in batch.items()
            }

            loss = model(**batch).loss
            loss.backward()
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)
            running += float(loss.item())

        model.eval()
        val_loss = 0.0
        n = 0
        with torch.no_grad():
            for batch in val_loader:
                batch = {
                    k: (v.to(device) if torch.is_tensor(v) else v)
                    for k, v in batch.items()
                }
                val_loss += float(model(**batch).loss.item())
                n += 1

        val_loss /= max(n, 1)
        print(f"epoch={epoch+1} validation_loss={val_loss:.4f}")

        if val_loss < best:
            best = val_loss
            model.save_pretrained(outdir)
            processor.save_pretrained(outdir)

    model.save_pretrained(outdir)
    processor.save_pretrained(outdir)

if __name__ == "__main__":
    main()