"""YOLOv11-seg instance segmentation for tuber detection.

Wraps the Ultralytics YOLO API with conf=0.25, imgsz=1280, default NMS
IoU, no augmentation. Polygon coordinates are normalized to [0, 1] and
written in the YOLO-seg .txt format consumed by io.read_yolo_seg_txt.
"""

from __future__ import annotations

import gc
from collections.abc import Iterable, Iterator
from pathlib import Path
from typing import Any

import numpy as np
from ultralytics import YOLO

from .io import IMG_EXTS

YOLO_PREDICT_DEFAULTS: dict[str, Any] = {
    "conf": 0.25,
    "imgsz": 1280,
    "verbose": False,
}


def load_model(weights_path: str | Path) -> YOLO:
    """Load YOLOv11-seg weights via the Ultralytics API."""
    return YOLO(str(weights_path))


def predict_image(
    model: YOLO,
    image: str | Path | np.ndarray,
    **predict_kwargs: Any,
) -> list[tuple[int, np.ndarray]]:
    """Run inference on a single image; return [(class_id, polygon_normalized)].

    `polygon_normalized` is an (N, 2) float array of (x, y) coordinates in
    [0, 1]. Returns an empty list if no instances were detected.
    """
    params = {**YOLO_PREDICT_DEFAULTS, **predict_kwargs}
    results = model.predict(source=image, stream=False, **params)
    r = results[0]
    if r.masks is None:
        return []
    h, w = r.orig_shape
    detections: list[tuple[int, np.ndarray]] = []
    for poly_xy, box in zip(r.masks.xy, r.boxes):
        class_id = int(box.cls[0])
        normalized = poly_xy / np.array([w, h], dtype=poly_xy.dtype)
        detections.append((class_id, normalized))
    return detections


def write_yolo_seg_txt(detections: Iterable[tuple[int, np.ndarray]], txt_path: str | Path) -> None:
    """Write detections in the legacy YOLO-seg format.

    One line per detection: `class_id x1 y1 x2 y2 ...` with coords in [0, 1].
    Creates an empty file if `detections` is empty (matches legacy behavior).
    """
    with open(txt_path, "w") as f:
        for class_id, poly in detections:
            coords = poly.reshape(-1)
            f.write(f"{class_id} " + " ".join(map(str, coords)) + "\n")


def _list_images(folder: Path) -> list[Path]:
    return sorted(p for p in folder.rglob("*") if p.suffix.lower() in IMG_EXTS)


def _chunked(lst: list[Path], n: int) -> Iterator[list[Path]]:
    for i in range(0, len(lst), n):
        yield lst[i : i + n]


def predict_directory(
    model: YOLO,
    input_dir: str | Path,
    output_dir: str | Path,
    chunk_size: int = 15,
    **predict_kwargs: Any,
) -> dict[str, int]:
    """Run inference over every image under input_dir, writing one .txt per image.

    Streams predictions in chunks of chunk_size and triggers GC between
    chunks. Empty .txt files are still created for images with no detections.

    Returns a dict with n_images, n_with_masks, n_seen for QC.
    """
    input_dir = Path(input_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    image_paths = _list_images(input_dir)
    print(f"[INFO] {input_dir.name}: {len(image_paths)} images")
    if not image_paths:
        return {"n_images": 0, "n_seen": 0, "n_with_masks": 0}

    params = {**YOLO_PREDICT_DEFAULTS, **predict_kwargs}
    seen = 0
    with_masks = 0

    for ci, part in enumerate(_chunked(image_paths, chunk_size), start=1):
        gen = model.predict(source=[p.as_posix() for p in part], stream=True, **params)
        # zip() rather than trusting r.path: in some Ultralytics versions
        # r.path falls back to "imageN" inside a batched/streamed call,
        # which causes chunks to overwrite each other's outputs. stream=True
        # preserves input order, so part[i] matches the i-th yielded result.
        for img_path, r in zip(part, gen):
            seen += 1
            txt_path = output_dir / f"{img_path.stem}.txt"
            with open(txt_path, "w") as f:
                if r.masks is not None:
                    with_masks += 1
                    h, w = r.orig_shape
                    for poly_xy, box in zip(r.masks.xy, r.boxes):
                        class_id = int(box.cls[0])
                        coords: list[float] = []
                        for x, y in poly_xy:
                            coords.append(x / w)
                            coords.append(y / h)
                        f.write(f"{class_id} " + " ".join(map(str, coords)) + "\n")
            del r
        gc.collect()
        if ci == 1 or ci % 10 == 0:
            print(
                f"[INFO] chunk {ci}: {min(ci * chunk_size, len(image_paths))}/"
                f"{len(image_paths)} images"
            )

    print(f"[DONE] {output_dir}: seen={seen}, with_masks={with_masks}")
    return {"n_images": len(image_paths), "n_seen": seen, "n_with_masks": with_masks}