#!/usr/bin/env python3
"""Quantify exact interior-object repetition caused by overlapping YOLO tiles."""

from __future__ import annotations

import argparse
import json
import re
from collections import defaultdict
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

try:
    from scripts.audit_yolo_label_relationships import (
        Box,
        parse_label_file,
        resolve_path,
        sha256_file,
    )
except ModuleNotFoundError:  # Standalone operator-tool copy beside the auditor.
    from audit_yolo_label_relationships import (
        Box,
        parse_label_file,
        resolve_path,
        sha256_file,
    )


OFFSET_PATTERN = re.compile(r"_r(?P<row>\d+)_c(?P<column>\d+)$")


def tile_offsets(image_path: str) -> tuple[int, int]:
    match = OFFSET_PATTERN.search(Path(image_path).stem)
    if match is None:
        raise ValueError(f"tile filename has no row/column offsets: {image_path}")
    return int(match["row"]), int(match["column"])


def interior_global_key(
    box: Box,
    *,
    row_offset: int,
    column_offset: int,
    tile_size: int,
    edge_tolerance_pixels: float,
) -> tuple[float, float, float, float] | None:
    left, top, right, bottom = box.coordinates
    pixel_box = (
        left * tile_size,
        top * tile_size,
        right * tile_size,
        bottom * tile_size,
    )
    if (
        pixel_box[0] <= edge_tolerance_pixels
        or pixel_box[1] <= edge_tolerance_pixels
        or tile_size - pixel_box[2] <= edge_tolerance_pixels
        or tile_size - pixel_box[3] <= edge_tolerance_pixels
    ):
        return None
    return (
        round(column_offset + pixel_box[0], 3),
        round(row_offset + pixel_box[1], 3),
        round(column_offset + pixel_box[2], 3),
        round(row_offset + pixel_box[3], 3),
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--summary-path", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--edge-tolerance-pixels", type=float, default=0.5)
    args = parser.parse_args()
    if args.output.exists():
        parser.error(f"output already exists: {args.output}")
    if args.edge_tolerance_pixels < 0:
        parser.error("--edge-tolerance-pixels must be non-negative")

    summary_path = args.summary_path.expanduser().resolve(strict=True)
    summary = json.loads(summary_path.read_text(encoding="utf-8"))
    tiles = summary.get("tiles")
    tile_size = summary.get("tile_size")
    if not isinstance(tiles, list):
        raise ValueError("dataset summary must contain a tiles list")
    if not isinstance(tile_size, int) or tile_size <= 0:
        raise ValueError("dataset summary must contain a positive integer tile_size")

    groups: dict[tuple[str, int, float, float, float, float], list[dict[str, Any]]] = (
        defaultdict(list)
    )
    reviewed_label_count = 0
    edge_label_count = 0
    for tile in tiles:
        if not isinstance(tile, dict) or not tile.get("kept", True):
            continue
        image_path = str(tile.get("image_path") or "")
        row_offset, column_offset = tile_offsets(image_path)
        label_path = resolve_path(str(tile.get("label_path") or ""), summary_path)
        boxes = parse_label_file(label_path)
        reviewed_label_count += len(boxes)
        for index, box in enumerate(boxes):
            global_key = interior_global_key(
                box,
                row_offset=row_offset,
                column_offset=column_offset,
                tile_size=tile_size,
                edge_tolerance_pixels=args.edge_tolerance_pixels,
            )
            if global_key is None:
                edge_label_count += 1
                continue
            key = (
                str(tile.get("sample_slug") or "unknown"),
                box.class_id,
                *global_key,
            )
            groups[key].append(
                {
                    "split": str(tile.get("split") or "unknown"),
                    "tile_index": int(tile.get("tile_index") or 0),
                    "label_index": index,
                    "image_path": image_path,
                    "label_path": str(label_path),
                }
            )

    repeated_groups = [members for members in groups.values() if len(members) > 1]
    split_leakage_groups = [
        members
        for members in repeated_groups
        if len({member["split"] for member in members}) > 1
    ]
    repetition_histogram: dict[str, int] = defaultdict(int)
    for members in groups.values():
        repetition_histogram[str(len(members))] += 1
    sample_stats: dict[str, dict[str, int]] = defaultdict(
        lambda: {
            "unique_interior_objects": 0,
            "interior_label_rows": 0,
            "repeated_object_groups": 0,
            "extra_repeated_rows": 0,
        }
    )
    for key, members in groups.items():
        sample = sample_stats[key[0]]
        sample["unique_interior_objects"] += 1
        sample["interior_label_rows"] += len(members)
        if len(members) > 1:
            sample["repeated_object_groups"] += 1
            sample["extra_repeated_rows"] += len(members) - 1

    interior_label_count = sum(len(members) for members in groups.values())
    extra_repeated_rows = sum(len(members) - 1 for members in repeated_groups)
    payload = {
        "schema_version": 1,
        "generated_at": datetime.now(UTC).isoformat(),
        "status": "attention" if repeated_groups else "ok",
        "claim_boundary": (
            "Exact reconstructed repetition among interior boxes only; tile-edge "
            "rows remain unlinked and repetition is not automatically an error."
        ),
        "summary_path": str(summary_path),
        "summary_sha256": sha256_file(summary_path),
        "tile_size": tile_size,
        "stride": summary.get("stride"),
        "edge_tolerance_pixels": args.edge_tolerance_pixels,
        "reviewed_label_count": reviewed_label_count,
        "interior_label_count": interior_label_count,
        "edge_label_count": edge_label_count,
        "unique_interior_object_count": len(groups),
        "repeated_object_group_count": len(repeated_groups),
        "extra_repeated_label_row_count": extra_repeated_rows,
        "max_repetition": max((len(members) for members in groups.values()), default=0),
        "repetition_histogram": dict(
            sorted(repetition_histogram.items(), key=lambda item: int(item[0]))
        ),
        "cross_split_repetition_group_count": len(split_leakage_groups),
        "sample_stats": [
            {"sample_slug": slug, **stats}
            for slug, stats in sorted(sample_stats.items())
        ],
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(json.dumps(payload, indent=2))
    return 0


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