#!/usr/bin/env python3
"""Compare local detection checkpoints on one non-protected YOLO validation split."""

from __future__ import annotations

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


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def safe_name(path: Path) -> str:
    value = re.sub(r"[^a-z0-9]+", "-", path.stem.casefold()).strip("-")
    return (value or path.stem.casefold())[:80]


def validation_images(dataset_yaml: Path) -> list[Path]:
    import yaml

    payload = yaml.safe_load(dataset_yaml.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError("dataset YAML must contain a mapping")
    source = payload.get("val")
    root = Path(str(payload.get("path") or dataset_yaml.parent)).expanduser()
    if not root.is_absolute():
        root = (dataset_yaml.parent / root).resolve(strict=False)
    if not isinstance(source, str) or not source.strip():
        raise ValueError("dataset YAML requires one explicit val image directory")
    directory = Path(source).expanduser()
    if not directory.is_absolute():
        directory = root / directory
    if not directory.is_dir():
        raise ValueError(f"validation image directory is unavailable: {directory}")
    images = sorted(
        path
        for path in directory.iterdir()
        if path.is_file()
        and path.suffix.casefold() in {".jpg", ".jpeg", ".png", ".tif", ".tiff"}
    )
    if not images:
        raise ValueError("validation image directory is empty")
    return images


def dataset_overlap_evidence(dataset_yaml: Path) -> dict[str, Any]:
    summary_path = dataset_yaml.parent / "yolo_tile_dataset_summary.json"
    if not summary_path.is_file():
        return {
            "status": "unavailable",
            "summary_path": str(summary_path),
            "validation_tiles_non_overlapping": None,
            "statistical_independence_established": False,
        }
    payload = json.loads(summary_path.read_text(encoding="utf-8"))
    tile_size = payload.get("tile_size")
    stride = payload.get("stride")
    if not isinstance(tile_size, int) or not isinstance(stride, int) or stride <= 0:
        return {
            "status": "invalid",
            "summary_path": str(summary_path),
            "summary_sha256": sha256_file(summary_path),
            "validation_tiles_non_overlapping": None,
            "statistical_independence_established": False,
        }
    overlap_pixels = max(tile_size - stride, 0)
    return {
        "status": "overlapping" if overlap_pixels else "non_overlapping",
        "summary_path": str(summary_path),
        "summary_sha256": sha256_file(summary_path),
        "tile_size": tile_size,
        "stride": stride,
        "overlap_pixels": overlap_pixels,
        "validation_tiles_non_overlapping": overlap_pixels == 0,
        "statistical_independence_established": False,
        "interpretation": (
            "Tile metrics can repeat the same source object and are valid for "
            "candidate ranking only, not independent object-level uncertainty."
            if overlap_pixels
            else "Tiles do not overlap, but spatial/statistical independence is "
            "not established because adjacent tiles share AOI context and edge "
            "objects can remain split."
        ),
    }


def _sample_slugs(payload: dict[str, Any], *, split: str | None) -> set[str]:
    tiles = payload.get("tiles")
    if not isinstance(tiles, list):
        raise ValueError("dataset summary requires a tiles list")
    samples: set[str] = set()
    for tile in tiles:
        if not isinstance(tile, dict):
            raise ValueError("dataset summary tiles must contain mappings")
        if split is not None and tile.get("split") != split:
            continue
        sample_slug = tile.get("sample_slug")
        if not isinstance(sample_slug, str) or not sample_slug.strip():
            raise ValueError("every selected tile requires a sample_slug")
        samples.add(sample_slug.strip())
    return samples


def model_lineage_independence_evidence(
    dataset_yaml: Path, lineage_summaries: list[Path]
) -> dict[str, Any]:
    evaluation_summary = dataset_yaml.parent / "yolo_tile_dataset_summary.json"
    if not evaluation_summary.is_file():
        return {
            "status": "unavailable",
            "reason": "evaluation dataset summary is unavailable",
            "evaluation_summary_path": str(evaluation_summary),
            "independent_for_all_supplied_lineage_corpora": False,
        }

    try:
        evaluation_payload = json.loads(evaluation_summary.read_text(encoding="utf-8"))
        evaluation_samples = _sample_slugs(evaluation_payload, split="val")
        if not evaluation_samples:
            raise ValueError("evaluation summary contains no validation samples")
    except (json.JSONDecodeError, OSError, ValueError) as exc:
        return {
            "status": "invalid",
            "reason": str(exc),
            "evaluation_summary_path": str(evaluation_summary),
            "evaluation_summary_sha256": sha256_file(evaluation_summary),
            "independent_for_all_supplied_lineage_corpora": False,
        }

    rows: list[dict[str, Any]] = []
    union_overlap: set[str] = set()
    for raw_summary in lineage_summaries:
        summary = raw_summary.expanduser().resolve(strict=True)
        try:
            payload = json.loads(summary.read_text(encoding="utf-8"))
            all_samples = _sample_slugs(payload, split=None)
            if not all_samples:
                raise ValueError("lineage summary contains no samples")
            roles_by_sample: dict[str, set[str]] = {}
            for tile in payload["tiles"]:
                sample_slug = str(tile["sample_slug"]).strip()
                split = str(tile.get("split") or "unknown").strip()
                roles_by_sample.setdefault(sample_slug, set()).add(split)
        except (json.JSONDecodeError, OSError, ValueError) as exc:
            return {
                "status": "invalid",
                "reason": f"{summary}: {exc}",
                "evaluation_summary_path": str(evaluation_summary),
                "evaluation_summary_sha256": sha256_file(evaluation_summary),
                "evaluation_samples": sorted(evaluation_samples),
                "independent_for_all_supplied_lineage_corpora": False,
            }
        overlap = evaluation_samples & all_samples
        union_overlap.update(overlap)
        rows.append(
            {
                "lineage_summary_path": str(summary),
                "lineage_summary_sha256": sha256_file(summary),
                "lineage_sample_count": len(all_samples),
                "overlapping_evaluation_samples": sorted(overlap),
                "exposure_roles": {
                    sample: sorted(roles_by_sample[sample]) for sample in sorted(overlap)
                },
            }
        )

    if not rows:
        return {
            "status": "unavailable",
            "reason": "no complete model-lineage summaries were supplied",
            "evaluation_summary_path": str(evaluation_summary),
            "evaluation_summary_sha256": sha256_file(evaluation_summary),
            "evaluation_samples": sorted(evaluation_samples),
            "lineage_corpora": [],
            "overlapping_evaluation_samples": [],
            "independent_for_all_supplied_lineage_corpora": False,
            "interpretation": (
                "Model-lineage/evaluation independence cannot be established "
                "without every ancestral corpus summary."
            ),
        }

    independent = not union_overlap
    return {
        "status": "independent" if independent else "overlap",
        "evaluation_summary_path": str(evaluation_summary),
        "evaluation_summary_sha256": sha256_file(evaluation_summary),
        "evaluation_samples": sorted(evaluation_samples),
        "lineage_corpora": rows,
        "overlapping_evaluation_samples": sorted(union_overlap),
        "independent_for_all_supplied_lineage_corpora": independent,
        "interpretation": (
            "No evaluation AOI occurs in any split of any supplied ancestral corpus."
            if independent
            else "At least one evaluation AOI was exposed in an ancestral train, "
            "validation, calibration or other split; the matrix is blocked before "
            "model loading."
        ),
    }


def training_sample_independence_evidence(
    dataset_yaml: Path, training_summaries: list[Path]
) -> dict[str, Any]:
    """Backward-compatible alias; evidence now checks every lineage split."""
    return model_lineage_independence_evidence(dataset_yaml, training_summaries)


def write_blocked_manifest(
    output: Path, dataset_yaml: Path, evidence: dict[str, Any]
) -> None:
    payload = {
        "schema_version": 3,
        "generated_at": datetime.now(UTC).isoformat(),
        "status": "blocked_model_lineage_sample_exposure",
        "claim_boundary": "No checkpoint ranking or release claim is permitted.",
        "dataset_yaml": str(dataset_yaml),
        "dataset_yaml_sha256": sha256_file(dataset_yaml),
        "model_lineage_independence_evidence": evidence,
        "model_loading_attempted": False,
        "gpu_inference_attempted": False,
    }
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )


def metric_value(metrics: Any, attribute: str) -> float:
    value = getattr(metrics.box, attribute)
    return float(value)


def background_detection_count(
    model: Any,
    images: list[Path],
    *,
    prefixes: tuple[str, ...],
    confidence: float,
    image_size: int,
    device: str,
) -> tuple[int, int]:
    selected = [path for path in images if path.stem.casefold().startswith(prefixes)]
    if not selected:
        raise ValueError(
            "no validation images match the declared pure-background prefixes"
        )
    count = 0
    for start in range(0, len(selected), 16):
        results = model.predict(
            [str(path) for path in selected[start : start + 16]],
            conf=confidence,
            iou=0.7,
            max_det=1000,
            imgsz=image_size,
            device=device,
            verbose=False,
        )
        count += sum(len(result.boxes) for result in results)
    return len(selected), count


def validate_pure_background_prefixes(
    images: list[Path], prefixes: tuple[str, ...]
) -> list[Path]:
    selected = [path for path in images if path.stem.casefold().startswith(prefixes)]
    if not selected:
        raise ValueError(
            "no validation images match the declared pure-background prefixes"
        )
    nonempty: list[str] = []
    for image in selected:
        try:
            relative = image.relative_to(image.parents[1])
        except ValueError as exc:
            raise ValueError(f"cannot resolve label path for {image}") from exc
        label = image.parents[1].parent / "labels" / relative.with_suffix(".txt")
        if not label.is_file():
            raise ValueError(f"pure-background image has no label file: {image}")
        if label.read_text(encoding="utf-8").strip():
            nonempty.append(str(label))
    if nonempty:
        raise ValueError(
            "pure-background prefixes include non-empty labels: "
            + ", ".join(nonempty[:10])
        )
    return selected


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset-yaml", type=Path, required=True)
    parser.add_argument("--model", type=Path, action="append", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--background-prefix", action="append", default=[])
    parser.add_argument("--background-confidence", type=float, default=0.15)
    parser.add_argument(
        "--lineage-summary",
        "--training-summary",
        dest="lineage_summary",
        type=Path,
        action="append",
        default=[],
        help="tile summary for every corpus in the complete model ancestry",
    )
    parser.add_argument(
        "--require-lineage-sample-independence",
        "--require-training-sample-independence",
        dest="require_lineage_sample_independence",
        action="store_true",
    )
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--batch", type=int, default=8)
    args = parser.parse_args()

    if args.output.exists():
        parser.error(f"output already exists: {args.output}")
    if not 0.0 < args.background_confidence < 1.0:
        parser.error("--background-confidence must be between zero and one")
    dataset_yaml = args.dataset_yaml.expanduser().resolve(strict=True)
    images = validation_images(dataset_yaml)
    overlap_evidence = dataset_overlap_evidence(dataset_yaml)
    independence_evidence = model_lineage_independence_evidence(
        dataset_yaml, args.lineage_summary
    )
    if args.require_lineage_sample_independence and not args.lineage_summary:
        parser.error(
            "--require-lineage-sample-independence requires every ancestral "
            "--lineage-summary"
        )
    if (
        args.require_lineage_sample_independence
        and not independence_evidence[
            "independent_for_all_supplied_lineage_corpora"
        ]
    ):
        write_blocked_manifest(args.output, dataset_yaml, independence_evidence)
        return 3
    prefixes = tuple(value.casefold() for value in args.background_prefix)
    if prefixes:
        validate_pure_background_prefixes(images, prefixes)

    import torch
    from ultralytics import YOLO

    if not torch.cuda.is_available() or not args.device.casefold().startswith("cuda"):
        raise SystemExit("checkpoint matrix requires the configured CUDA device")

    output_root = args.output.parent / f"{args.output.stem}-runs"
    rows: list[dict[str, Any]] = []
    seen_hashes: set[str] = set()
    for raw_model in args.model:
        model_path = raw_model.expanduser().resolve(strict=True)
        model_sha256 = sha256_file(model_path)
        if model_sha256 in seen_hashes:
            continue
        seen_hashes.add(model_sha256)
        row: dict[str, Any] = {
            "model_path": str(model_path),
            "model_sha256": model_sha256,
            "size_bytes": model_path.stat().st_size,
        }
        model = None
        try:
            model = YOLO(str(model_path))
            metrics = model.val(
                data=str(dataset_yaml),
                split="val",
                imgsz=args.imgsz,
                batch=args.batch,
                workers=0,
                device=args.device,
                project=str(output_root),
                name=safe_name(model_path),
                exist_ok=False,
                plots=False,
                save_json=False,
                verbose=False,
            )
            row.update(
                {
                    "status": "ok",
                    "precision": metric_value(metrics, "mp"),
                    "recall": metric_value(metrics, "mr"),
                    "map50": metric_value(metrics, "map50"),
                    "map50_95": metric_value(metrics, "map"),
                }
            )
            if prefixes:
                background_images, background_detections = background_detection_count(
                    model,
                    images,
                    prefixes=prefixes,
                    confidence=args.background_confidence,
                    image_size=args.imgsz,
                    device=args.device,
                )
                row.update(
                    {
                        "pure_background_image_count": background_images,
                        "pure_background_detection_count": background_detections,
                    }
                )
        except Exception as exc:  # preserve the complete attempted matrix
            row.update(
                {
                    "status": "error",
                    "error_type": type(exc).__name__,
                    "error": str(exc)[:1000],
                }
            )
        finally:
            del model
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        rows.append(row)
        print(json.dumps(row, sort_keys=True), flush=True)

    successful = [row for row in rows if row["status"] == "ok"]
    successful.sort(
        key=lambda row: (
            int(row.get("pure_background_detection_count") == 0 and bool(prefixes)),
            row["map50_95"],
            row["map50"],
            row["precision"],
            row["recall"],
        ),
        reverse=True,
    )
    payload = {
        "schema_version": 3,
        "generated_at": datetime.now(UTC).isoformat(),
        "status": "ok" if successful else "failed",
        "claim_boundary": "Non-protected validation ranking only; no test, challenge or promotion claim.",
        "dataset_yaml": str(dataset_yaml),
        "dataset_yaml_sha256": sha256_file(dataset_yaml),
        "dataset_overlap_evidence": overlap_evidence,
        "model_lineage_independence_evidence": independence_evidence,
        "validation_image_count": len(images),
        "pure_background_prefixes": list(prefixes),
        "pure_background_confidence": args.background_confidence,
        "device": args.device,
        "torch_version": torch.__version__,
        "candidate_count": len(rows),
        "successful_candidate_count": len(successful),
        "ranking": successful,
        "attempts": rows,
    }
    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"
    )
    return 0 if successful else 2


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