from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
from torch.utils.data import DataLoader

from .boxes import box_cxcywh_to_xyxy
from .config import apply_overrides, load_config
from .data import build_dataset, detection_collate
from .model import build_model


@torch.inference_mode()
def evaluate_coco(
    model,
    data_loader,
    device: torch.device,
    output_file: str | Path,
    confidence: float = 0.001,
    max_detections: int = 300,
) -> dict[str, float]:
    try:
        from pycocotools.coco import COCO
        from pycocotools.cocoeval import COCOeval
    except ImportError as error:
        raise RuntimeError("COCO evaluation requires: pip install -e '.[coco]'") from error

    model.eval()
    input_size = model.spec.input_size
    predictions = []
    label_to_category = data_loader.dataset.label_to_category
    for images, targets in data_loader:
        images = images.to(device, non_blocking=True)
        outputs = model(images)
        probabilities = outputs["pred_logits"].sigmoid()
        normalized_boxes = box_cxcywh_to_xyxy(outputs["pred_boxes"]).clamp(0.0, 1.0)
        for batch_index, target in enumerate(targets):
            scores, labels = probabilities[batch_index].max(dim=-1)
            count = min(max_detections, scores.numel())
            scores, indices = scores.topk(count)
            keep = scores >= confidence
            scores, indices = scores[keep], indices[keep]
            labels = labels[indices]
            boxes = normalized_boxes[batch_index, indices] * input_size
            ratio, offset_x, offset_y = target["transform"].tolist()
            original_height, original_width = target["original_size"].tolist()
            boxes[:, [0, 2]] = (boxes[:, [0, 2]] - offset_x) / ratio
            boxes[:, [1, 3]] = (boxes[:, [1, 3]] - offset_y) / ratio
            boxes[:, [0, 2]].clamp_(0, original_width)
            boxes[:, [1, 3]].clamp_(0, original_height)
            boxes[:, 2:] -= boxes[:, :2]
            for box, score, label in zip(boxes.cpu(), scores.cpu(), labels.cpu(), strict=True):
                predictions.append(
                    {
                        "image_id": int(target["image_id"]),
                        "category_id": int(label_to_category[int(label)]),
                        "bbox": [round(float(value), 3) for value in box],
                        "score": float(score),
                    }
                )

    output_file = Path(output_file)
    output_file.parent.mkdir(parents=True, exist_ok=True)
    with output_file.open("w", encoding="utf-8") as handle:
        json.dump(predictions, handle)
    ground_truth = COCO(str(data_loader.dataset.annotation_file))
    if not predictions:
        return {"AP": 0.0, "AP50": 0.0, "AP75": 0.0, "APS": 0.0, "APM": 0.0, "APL": 0.0}
    detections = ground_truth.loadRes(str(output_file))
    evaluator = COCOeval(ground_truth, detections, "bbox")
    evaluator.params.imgIds = [int(item["id"]) for item in data_loader.dataset.images]
    evaluator.evaluate()
    evaluator.accumulate()
    evaluator.summarize()
    names = ("AP", "AP50", "AP75", "APS", "APM", "APL")
    return {name: float(evaluator.stats[index]) for index, name in enumerate(names)}


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate ObjectModel-v1 on COCO")
    parser.add_argument("--config", default="configs/objectmodel_v1.yaml")
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--data-root", required=True)
    parser.add_argument("--output", default="outputs/predictions.json")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--set", action="append", default=[])
    args = parser.parse_args()
    config = apply_overrides(load_config(args.config), args.set)
    model = build_model(config)
    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    model.load_state_dict(checkpoint.get("ema", checkpoint.get("model", checkpoint)))
    device = torch.device(args.device)
    model.to(device)
    dataset = build_dataset(config, args.data_root, "val")
    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.workers,
        pin_memory=device.type == "cuda",
        collate_fn=detection_collate,
    )
    print(json.dumps(evaluate_coco(model, loader, device, args.output), indent=2))


if __name__ == "__main__":
    main()