"""Run the first GeoAI object-detection baseline on drone images.

This baseline uses the opengeos/geoai environment plus a general-purpose
Ultralytics model. It intentionally keeps georeferenced GeoTIFF processing for
the next stage, because the current JPEGs do not carry usable GPS metadata.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import tempfile
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from PIL import Image


IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"}
COCO_TARGET_CLASS_IDS = [0, 1, 2, 3, 5, 7]


def parse_args() -> argparse.Namespace:
    root = Path(__file__).resolve().parents[2]
    default_input = root / "shared" / "data" / "raw" / "01-object-detection"
    default_output = root / "shared" / "outputs" / "01-object-detection"
    parser = argparse.ArgumentParser(description="Detect objects in drone images.")
    parser.add_argument("--input", type=Path, default=default_input)
    parser.add_argument("--output", type=Path, default=default_output)
    parser.add_argument("--model", default="yolo11n.pt", help="Ultralytics model name or path.")
    parser.add_argument("--confidence", type=float, default=0.20)
    parser.add_argument("--image-size", type=int, default=1024, help="Model input size for each tile.")
    parser.add_argument("--tile-size", type=int, default=1024, help="Pixel size of the sliding detection window.")
    parser.add_argument("--tile-overlap", type=float, default=0.20, help="Overlap ratio between adjacent windows.")
    return parser.parse_args()


def load_image_paths(input_dir: Path) -> tuple[list[Path], list[dict[str, str]]]:
    paths: list[Path] = []
    skipped: list[dict[str, str]] = []
    candidates = [input_dir] if input_dir.is_file() else sorted(input_dir.iterdir())
    for path in candidates:
        if not path.is_file() or path.suffix.lower() not in IMAGE_SUFFIXES:
            continue
        try:
            with Image.open(path) as image:
                image.verify()
                if image.format == "MPO":
                    skipped.append({"file": path.name, "reason": "MPO is not supported in baseline"})
                    continue
        except Exception as exc:  # pragma: no cover - depends on source files
            skipped.append({"file": path.name, "reason": f"unreadable: {exc}"})
            continue
        paths.append(path)
    return paths, skipped


def tile_starts(length: int, tile_size: int, overlap: float) -> list[int]:
    if not 0 <= overlap < 1:
        raise ValueError("tile-overlap must be between 0 and 1")
    if length <= tile_size:
        return [0]
    stride = max(1, int(tile_size * (1 - overlap)))
    starts = list(range(0, length - tile_size + 1, stride))
    last = length - tile_size
    if starts[-1] != last:
        starts.append(last)
    return starts


def intersection_over_union(box_a: list[float], box_b: list[float]) -> float:
    left = max(box_a[0], box_b[0])
    top = max(box_a[1], box_b[1])
    right = min(box_a[2], box_b[2])
    bottom = min(box_a[3], box_b[3])
    intersection = max(0.0, right - left) * max(0.0, bottom - top)
    area_a = max(0.0, box_a[2] - box_a[0]) * max(0.0, box_a[3] - box_a[1])
    area_b = max(0.0, box_b[2] - box_b[0]) * max(0.0, box_b[3] - box_b[1])
    union = area_a + area_b - intersection
    return intersection / union if union else 0.0


def intersection_over_smaller(box_a: list[float], box_b: list[float]) -> float:
    left = max(box_a[0], box_b[0])
    top = max(box_a[1], box_b[1])
    right = min(box_a[2], box_b[2])
    bottom = min(box_a[3], box_b[3])
    intersection = max(0.0, right - left) * max(0.0, bottom - top)
    area_a = max(0.0, box_a[2] - box_a[0]) * max(0.0, box_a[3] - box_a[1])
    area_b = max(0.0, box_b[2] - box_b[0]) * max(0.0, box_b[3] - box_b[1])
    smaller_area = min(area_a, area_b)
    return intersection / smaller_area if smaller_area else 0.0


def class_aware_nms(
    detections: list[dict[str, Any]],
    iou_threshold: float = 0.50,
    containment_threshold: float = 0.80,
) -> list[dict[str, Any]]:
    kept: list[dict[str, Any]] = []
    for candidate in sorted(detections, key=lambda item: item["confidence"], reverse=True):
        if all(
            candidate["class_id"] != existing["class_id"]
            or (
                intersection_over_union(candidate["bbox_xyxy"], existing["bbox_xyxy"]) < iou_threshold
                and intersection_over_smaller(candidate["bbox_xyxy"], existing["bbox_xyxy"])
                < containment_threshold
            )
            for existing in kept
        ):
            kept.append(candidate)
    return kept


def predict_tiled(model: Any, image_path: Path, args: argparse.Namespace, np: Any) -> tuple[list[dict[str, Any]], int, int]:
    with Image.open(image_path) as source:
        image = source.convert("RGB")
    width, height = image.size
    detections: list[dict[str, Any]] = []
    names: dict[int, str] = {}
    for y0 in tile_starts(height, args.tile_size, args.tile_overlap):
        for x0 in tile_starts(width, args.tile_size, args.tile_overlap):
            tile = image.crop((x0, y0, min(x0 + args.tile_size, width), min(y0 + args.tile_size, height)))
            result = model.predict(
                source=np.asarray(tile),
                device="cpu",
                imgsz=args.image_size,
                conf=args.confidence,
                classes=COCO_TARGET_CLASS_IDS,
                max_det=100,
                verbose=False,
            )[0]
            names = result.names
            if result.boxes is None:
                continue
            for box in result.boxes:
                class_id = int(box.cls.item())
                local_box = [float(value) for value in box.xyxy[0].tolist()]
                detections.append(
                    {
                        "class_id": class_id,
                        "class_name": names[class_id],
                        "confidence": round(float(box.conf.item()), 4),
                        "bbox_xyxy": [
                            round(local_box[0] + x0, 2),
                            round(local_box[1] + y0, 2),
                            round(local_box[2] + x0, 2),
                            round(local_box[3] + y0, 2),
                        ],
                    }
                )
    return class_aware_nms(detections), width, height


def draw_detections(image_path: Path, detections: list[dict[str, Any]], output_path: Path) -> None:
    import cv2
    import numpy as np

    canvas = cv2.cvtColor(np.asarray(Image.open(image_path).convert("RGB")), cv2.COLOR_RGB2BGR)
    for detection in detections:
        x1, y1, x2, y2 = [int(round(value)) for value in detection["bbox_xyxy"]]
        label = f"{detection['class_name']} {detection['confidence']:.2f}"
        cv2.rectangle(canvas, (x1, y1), (x2, y2), (255, 80, 0), 4)
        cv2.putText(canvas, label, (x1, max(30, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 80, 0), 2)
    cv2.imwrite(str(output_path), canvas)


def main() -> int:
    args = parse_args()
    input_dir = args.input.resolve()
    output_dir = args.output.resolve()
    annotated_dir = output_dir / "annotated"
    annotated_dir.mkdir(parents=True, exist_ok=True)
    cache_root = Path(tempfile.gettempdir()) / "geoai-object-detection"
    os.environ.setdefault("YOLO_CONFIG_DIR", str(cache_root / "ultralytics"))
    os.environ.setdefault("MPLCONFIGDIR", str(cache_root / "matplotlib"))
    if not input_dir.exists():
        print(f"Input directory does not exist: {input_dir}", file=sys.stderr)
        return 2

    image_paths, skipped = load_image_paths(input_dir)
    if not image_paths:
        print("No supported images found.", file=sys.stderr)
        return 2

    # Imports happen after argument validation so the CLI can explain path
    # mistakes without requiring the heavyweight ML stack.
    import geoai
    import numpy as np
    import torch
    import ultralytics
    from ultralytics import YOLO

    started = time.perf_counter()
    model = YOLO(args.model)
    detections: list[dict[str, Any]] = []
    for image_path in image_paths:
        image_detections, width, height = predict_tiled(model, image_path, args, np)
        annotated_path = annotated_dir / image_path.name
        draw_detections(image_path, image_detections, annotated_path)
        detections.append(
            {
                "file": image_path.name,
                "annotated_file": str(annotated_path.relative_to(output_dir)),
                "width": width,
                "height": height,
                "detections": image_detections,
            }
        )

    elapsed = round(time.perf_counter() - started, 3)
    (output_dir / "detections.json").write_text(
        json.dumps({"images": detections}, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    metadata = {
        "created_at": datetime.now(UTC).isoformat(),
        "geoai_package": getattr(geoai, "__version__", "unknown"),
        "ultralytics": ultralytics.__version__,
        "torch": torch.__version__,
        "device": "cpu",
        "cuda_available": bool(torch.cuda.is_available()),
        "model": args.model,
        "confidence": args.confidence,
        "image_size": args.image_size,
        "tile_size": args.tile_size,
        "tile_overlap": args.tile_overlap,
        "merge_iou_threshold": 0.50,
        "merge_containment_threshold": 0.80,
        "target_class_ids": COCO_TARGET_CLASS_IDS,
        "target_class_note": "COCO person, bicycle, car, motorcycle, bus and truck; tree is not a COCO class.",
        "input_dir": str(input_dir),
        "processed_images": len(image_paths),
        "skipped_images": skipped,
        "detection_count": sum(len(item["detections"]) for item in detections),
        "elapsed_seconds": elapsed,
        "notes": [
            "This is a baseline for people and common vehicle classes.",
            "The default COCO model does not provide a tree class.",
            "JPEG inputs have no usable GPS metadata; GeoJSON is deferred to GeoTIFF stage.",
        ],
    }
    (output_dir / "run_metadata.json").write_text(
        json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(f"Processed {len(image_paths)} images; detections={metadata['detection_count']}; elapsed={elapsed}s")
    print(f"Annotated images: {annotated_dir}")
    print(f"JSON: {output_dir / 'detections.json'}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())