"""
Fit YOLO-style anchor boxes using k-means based on COCO-format annotations.

Generated by gpt-5.2-codex xhigh.
"""

import argparse
import json
import logging
import math
from pathlib import Path
from pprint import pformat
from typing import Any
from typing import Literal
from typing import Optional
from typing import TypedDict

import torch

from birder.common import cli
from birder.conf import settings

logger = logging.getLogger(__name__)


class AnchorPreset(TypedDict):
    num_scales: int
    num_anchors: int
    default_size: tuple[int, int]
    format: Literal["pixels", "grid"]


MODEL_PRESETS: dict[str, AnchorPreset] = {
    "yolo_v2": {"num_scales": 1, "num_anchors": 5, "default_size": (416, 416), "format": "grid"},
    "yolo_v3": {"num_scales": 3, "num_anchors": 9, "default_size": (416, 416), "format": "pixels"},
    "yolo_v4": {"num_scales": 3, "num_anchors": 9, "default_size": (608, 608), "format": "pixels"},
    "yolo_v4_tiny": {"num_scales": 2, "num_anchors": 6, "default_size": (416, 416), "format": "pixels"},
}


def _load_ignore_list(ignore_file: Optional[str]) -> set[str]:
    if ignore_file is None:
        return set()

    with open(ignore_file, "r", encoding="utf-8") as handle:
        return {line.strip() for line in handle if line.strip()}


def _load_coco_boxes(
    coco_json_path: str, target_size: tuple[int, int], ignore_list: set[str], min_size: float, ignore_crowd: bool
) -> tuple[torch.Tensor, dict[str, int]]:
    coco_path = Path(coco_json_path)
    if coco_path.exists() is False:
        raise ValueError(f"COCO json not found at {coco_path}")

    with open(coco_path, "r", encoding="utf-8") as handle:
        data = json.load(handle)

    images = {}
    for image in data.get("images", []):
        image_id = image.get("id")
        if image_id is None:
            continue
        images[image_id] = (image.get("width"), image.get("height"), image.get("file_name", ""))

    stats = {
        "total_annotations": 0,
        "used_annotations": 0,
        "crowd_annotations": 0,
        "invalid_bbox": 0,
        "ignored_images": 0,
        "missing_images": 0,
        "missing_size": 0,
        "too_small": 0,
    }
    boxes: list[tuple[float, float]] = []
    target_h = float(target_size[0])
    target_w = float(target_size[1])
    for annotation in data.get("annotations", []):
        stats["total_annotations"] += 1
        if ignore_crowd is True and annotation.get("iscrowd", 0) == 1:
            stats["crowd_annotations"] += 1
            continue

        bbox = annotation.get("bbox")
        if bbox is None or len(bbox) != 4:
            stats["invalid_bbox"] += 1
            continue

        image_id = annotation.get("image_id")
        if image_id not in images:
            stats["missing_images"] += 1
            continue

        img_w, img_h, file_name = images[image_id]
        if file_name in ignore_list:
            stats["ignored_images"] += 1
            continue

        if img_w in {None, 0} or img_h in {None, 0}:
            stats["missing_size"] += 1
            continue

        bbox_w = float(bbox[2])
        bbox_h = float(bbox[3])
        if bbox_w <= 0.0 or bbox_h <= 0.0:
            stats["invalid_bbox"] += 1
            continue

        scaled_w = bbox_w / float(img_w) * target_w
        scaled_h = bbox_h / float(img_h) * target_h
        if scaled_w < min_size or scaled_h < min_size:
            stats["too_small"] += 1
            continue

        boxes.append((scaled_w, scaled_h))

    stats["used_annotations"] = len(boxes)
    if len(boxes) == 0:
        raise ValueError("No valid bounding boxes found for anchor fitting")

    return (torch.tensor(boxes, dtype=torch.float32), stats)


def _wh_iou(boxes: torch.Tensor, anchors: torch.Tensor) -> torch.Tensor:
    boxes = boxes[:, None, :]
    anchors = anchors[None, :, :]
    inter = torch.min(boxes, anchors).prod(dim=2)
    union = boxes.prod(dim=2) + anchors.prod(dim=2) - inter
    return inter / (union + 1e-9)


def _kmeans_plusplus_init(boxes: torch.Tensor, num_anchors: int, generator: torch.Generator) -> torch.Tensor:
    n_local_trials = 2 + int(math.log(num_anchors))

    anchors = []

    first_idx = torch.randint(0, boxes.size(0), (1,), generator=generator).item()
    anchors.append(boxes[first_idx])
    for _ in range(num_anchors - 1):
        anchors_tensor = torch.stack(anchors)
        ious = _wh_iou(boxes, anchors_tensor)
        max_ious = ious.max(dim=1).values

        min_distances = 1.0 - max_ious
        squared_distances = min_distances**2

        probs = squared_distances / (squared_distances.sum() + 1e-9)
        cumulative_probs = torch.cumsum(probs, dim=0)
        r = torch.rand(n_local_trials, generator=generator)
        candidate_indices = torch.searchsorted(cumulative_probs, r, right=True)

        candidate_indices = torch.clamp(candidate_indices, 0, boxes.size(0) - 1)
        candidate_boxes = boxes[candidate_indices]
        candidate_ious = _wh_iou(boxes, candidate_boxes)  # (n_boxes, n_trials)
        candidate_distances = 1.0 - candidate_ious

        min_distances_expanded = min_distances.unsqueeze(1)
        candidate_potentials = torch.min(min_distances_expanded, candidate_distances).sum(dim=0)

        best_trial_idx = torch.argmin(candidate_potentials).item()
        best_idx = candidate_indices[best_trial_idx].item()
        anchors.append(boxes[best_idx])

    return torch.stack(anchors)


def _kmeans_anchors(
    boxes: torch.Tensor, num_anchors: int, seed: Optional[int], max_iter: int
) -> tuple[torch.Tensor, torch.Tensor]:
    if boxes.size(0) < num_anchors:
        raise ValueError(
            f"Not enough boxes ({boxes.size(0)}) to fit {num_anchors} anchors, Reduce --num-anchors or add more data"
        )

    generator = torch.Generator()
    if seed is not None:
        generator.manual_seed(seed)

    anchors = _kmeans_plusplus_init(boxes, num_anchors, generator)
    assignments = torch.full((boxes.size(0),), -1, dtype=torch.int64)

    for _ in range(max_iter):
        ious = _wh_iou(boxes, anchors)
        new_assignments = torch.argmax(ious, dim=1)
        if torch.equal(assignments, new_assignments):
            break

        assignments = new_assignments
        for idx in range(num_anchors):
            mask = assignments == idx
            if mask.any():
                anchors[idx] = boxes[mask].median(dim=0).values
            else:
                rand_idx = torch.randint(0, boxes.size(0), (1,), generator=generator).item()
                anchors[idx] = boxes[rand_idx]

    return (anchors, assignments)


def _format_anchor_groups(anchor_groups: list[torch.Tensor], precision: int) -> list[list[tuple[float, float]]]:
    formatted: list[list[tuple[float, float]]] = []
    for group in anchor_groups:
        formatted.append([(round(float(anchor[0]), precision), round(float(anchor[1]), precision)) for anchor in group])

    return formatted


def _validate_args(
    args: argparse.Namespace,
) -> tuple[tuple[int, int], int, int, Literal["pixels", "grid"], list[float]]:
    preset = MODEL_PRESETS.get(args.preset) if args.preset is not None else None
    size = cli.parse_size(args.size) if args.size is not None else (preset["default_size"] if preset else None)
    if size is None:
        raise cli.ValidationError("Missing --size. Provide --size or use a --preset")

    num_scales = args.num_scales if args.num_scales is not None else (preset["num_scales"] if preset else None)
    num_anchors = args.num_anchors if args.num_anchors is not None else (preset["num_anchors"] if preset else None)
    output_format = args.format if args.format is not None else (preset["format"] if preset else None)
    if num_scales is None or num_anchors is None or output_format is None:
        raise cli.ValidationError(
            "Missing configuration. Provide --num-scales, --num-anchors and --format or use a --preset"
        )
    if num_scales < 1:
        raise cli.ValidationError("--num-scales must be >= 1")
    if num_anchors < 1:
        raise cli.ValidationError("--num-anchors must be >= 1")
    if num_anchors % num_scales != 0:
        raise cli.ValidationError("--num-anchors must be divisible by --num-scales")

    strides: list[float] = []
    if output_format == "grid":
        if args.stride is None:
            raise cli.ValidationError("--format grid requires --stride values per scale")

        strides = [float(value) for value in args.stride]
        if len(strides) != num_scales:
            raise cli.ValidationError("--stride must provide one value per scale when --format grid is used")
        if any(value <= 0 for value in strides):
            raise cli.ValidationError("--stride values must be > 0")

    return (size, num_scales, num_anchors, output_format, strides)


# pylint: disable=too-many-locals
def auto_anchors(args: argparse.Namespace) -> None:
    size, num_scales, num_anchors, output_format, strides = _validate_args(args)

    ignore_list = _load_ignore_list(args.ignore_file)
    boxes, stats = _load_coco_boxes(
        args.coco_json_path, size, ignore_list, args.min_size, ignore_crowd=not args.include_crowd
    )

    if args.preset is not None:
        logger.info(f"Using preset {args.preset}")

    logger.info(f"Fitting anchors using size={size[0]}x{size[1]}")
    logger.info(
        f"Annotations: total={stats['total_annotations']}, used={stats['used_annotations']}, "
        f"crowd={stats['crowd_annotations']}, invalid={stats['invalid_bbox']}, "
        f"ignored={stats['ignored_images']}, missing={stats['missing_images']}, "
        f"missing_size={stats['missing_size']}, too_small={stats['too_small']}"
    )

    anchors, _assignments = _kmeans_anchors(boxes, num_anchors, args.seed, args.max_iter)
    areas = anchors.prod(dim=1)
    anchors = anchors[torch.argsort(areas)]
    anchors_per_scale = num_anchors // num_scales
    anchor_groups = [anchors[i : i + anchors_per_scale] for i in range(0, num_anchors, anchors_per_scale)]

    ious = _wh_iou(boxes, anchors)
    best_iou = ious.max(dim=1).values
    logger.info(f"Mean IoU: {best_iou.mean().item():.4f}")

    formatted_groups = _format_anchor_groups(anchor_groups, args.precision)
    anchors_output = None
    if output_format == "pixels":
        if num_scales == 1:
            formatted_anchors: Any = formatted_groups[0]
        else:
            formatted_anchors = formatted_groups

        print("Anchors (pixels):")
        print(pformat(formatted_anchors))
        anchors_output = formatted_anchors

    if output_format == "grid":
        grid_groups: list[torch.Tensor] = []
        for group, stride in zip(anchor_groups, strides):
            grid_group = group.clone()
            grid_group[:, 0] = grid_group[:, 0] / stride
            grid_group[:, 1] = grid_group[:, 1] / stride
            grid_groups.append(grid_group)

        formatted_grid = _format_anchor_groups(grid_groups, args.precision)
        if num_scales == 1:
            formatted_grid_output: Any = formatted_grid[0]
        else:
            formatted_grid_output = formatted_grid

        print("Anchors (grid units):")
        print(pformat(formatted_grid_output))
        anchors_output = formatted_grid_output

    if args.output is not None:
        payload = {
            "anchors": anchors_output,
            "format": output_format,
            "size": [size[0], size[1]],
        }
        if output_format == "grid":
            payload["strides"] = strides

        with open(args.output, "w", encoding="utf-8") as handle:
            json.dump(payload, handle, indent=2)

        logger.info(f"Wrote anchors to {args.output}")


def set_parser(subparsers: Any) -> None:
    subparser = subparsers.add_parser(
        "auto-anchors",
        allow_abbrev=False,
        help="fit YOLO anchors with k-means on a COCO dataset",
        description="fit YOLO anchors with k-means on a COCO dataset",
        epilog=(
            "Usage examples:\n"
            "python -m birder.tools auto-anchors --preset yolo_v4 --size 640 "
            "--coco-json-path data/detection_data/training_annotations_coco.json\n"
            "python -m birder.tools auto-anchors --size 640 --num-anchors 9 --num-scales 3 --format pixels "
            "--coco-json-path data/detection_data/training_annotations_coco.json\n"
            "python -m birder.tools auto-anchors --preset yolo_v4_tiny --size 416 416 "
            "--coco-json-path ~/Datasets/cocodataset/annotations/instances_train2017.json --output anchors.json\n"
            "python -m birder.tools auto-anchors --preset yolo_v2 --stride 32 "
            "--coco-json-path data/detection_data/training_annotations_coco.json\n"
            "python -m birder.tools auto-anchors --size 640 --num-anchors 9 --num-scales 3 "
            "--format grid --stride 8 16 32 --coco-json-path data/detection_data/training_annotations_coco.json\n"
        ),
        formatter_class=cli.ArgumentHelpFormatter,
    )
    subparser.add_argument(
        "--preset", type=str, choices=sorted(MODEL_PRESETS.keys()), help="YOLO preset for anchor formatting"
    )
    subparser.add_argument(
        "--size",
        type=int,
        nargs="+",
        metavar=("H", "W"),
        help="target image size as [height, width], required without --preset",
    )
    subparser.add_argument("--num-anchors", type=int, help="number of anchors to fit, required without --preset")
    subparser.add_argument("--num-scales", type=int, help="number of output scales, required without --preset")
    subparser.add_argument(
        "--format", type=str, choices=["pixels", "grid"], help="anchor output format, required without --preset"
    )
    subparser.add_argument(
        "--stride", type=int, nargs="+", default=[32], help="strides per scale used to convert anchors to grid units"
    )
    subparser.add_argument(
        "--min-size", type=float, default=1.0, help="minimum scaled box size to include in anchor fitting"
    )
    subparser.add_argument("--include-crowd", default=False, action="store_true", help="include crowd annotations")
    subparser.add_argument("--seed", type=int, help="random seed for k-means initialization")
    subparser.add_argument("--max-iter", type=int, default=1000, help="maximum k-means iterations")
    subparser.add_argument("--precision", type=int, default=1, help="number of decimals to keep in anchor output")
    subparser.add_argument(
        "--ignore-file", type=str, metavar="FILE", help="file containing image names to skip (one per line)"
    )
    subparser.add_argument(
        "--coco-json-path",
        type=str,
        default=f"{settings.TRAINING_DETECTION_ANNOTATIONS_PATH}_coco.json",
        help="training COCO json path",
    )
    subparser.add_argument("--output", type=str, help="write anchors as JSON to this path")
    subparser.set_defaults(func=main)


def main(args: argparse.Namespace) -> None:
    auto_anchors(args)
