#!/usr/bin/env python3
"""Collect audited seed-42/43/44 official-test results for the final manuscript."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
from pathlib import Path


METHODS = {
    "coco_full": {"label": "COCO-full", "multiplier": 1},
    "imagenet_coco": {"label": "ImageNet+COCO", "multiplier": 5},
    "simclr_coco": {"label": "SimCLR+COCO", "multiplier": 5},
    "supcon_coco": {"label": "SupCon+COCO", "multiplier": 5},
}
EXPECTED_TEST_ANNOTATION_SHA256 = (
    "9eb5805ca031cebcd70bab4a4aa0c89d98646c206d333803a8d5b10ea7fe9960"
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--old-eval-root", type=Path, required=True)
    parser.add_argument("--seed44-root", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    return parser.parse_args()


def sha256(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 read_json(path: Path) -> dict:
    if not path.is_file():
        raise FileNotFoundError(path)
    return json.loads(path.read_text(encoding="utf-8"))


def read_csv(path: Path) -> list[dict[str, str]]:
    if not path.is_file():
        raise FileNotFoundError(path)
    with path.open(encoding="utf-8-sig", newline="") as handle:
        return list(csv.DictReader(handle))


def mean_sample_sd(values: list[float]) -> tuple[float, float]:
    if not values:
        raise ValueError("Cannot summarize an empty sample")
    mean = sum(values) / len(values)
    if len(values) == 1:
        return mean, 0.0
    variance = sum((value - mean) ** 2 for value in values) / (len(values) - 1)
    return mean, math.sqrt(variance)


def select_all_3000(rows: list[dict[str, str]], iou_type: str) -> dict[str, float]:
    matches = [
        row
        for row in rows
        if row.get("iou_type") == iou_type
        and row.get("density") == "all"
        and int(row.get("max_dets", -1)) == 3000
    ]
    if len(matches) != 1:
        raise ValueError(
            f"Expected one {iou_type}/all/maxDets=3000 row, found {len(matches)}"
        )
    selected = matches[0]
    metrics = {key: float(selected[key]) for key in ("AP", "AP50", "AR")}
    if not all(math.isfinite(value) for value in metrics.values()):
        raise ValueError(f"Non-finite metrics in selected row: {selected}")
    if int(selected["num_images"]) != 1564:
        raise ValueError(f"Unexpected official-test image count: {selected['num_images']}")
    return metrics


def validate_protocol_evidence(
    *,
    test_dir: Path,
    method: str,
    seed: int,
    require_prediction_files: bool = True,
) -> tuple[dict[str, float], dict[str, float], dict]:
    protocol_dir = test_dir / "protocol_ablation"
    csv_path = protocol_dir / f"{method}_seed{seed}_dense_protocol_ablation.csv"
    audit_path = protocol_dir / f"{method}_seed{seed}_dense_protocol_ablation_audit.json"
    bbox_path = test_dir / "predictions_bbox.json"
    segm_path = test_dir / "predictions_segm.json"
    required_paths = [csv_path, audit_path]
    if require_prediction_files:
        required_paths.extend((bbox_path, segm_path))
    for path in required_paths:
        if not path.is_file() or path.stat().st_size == 0:
            raise FileNotFoundError(f"Missing or empty evidence file: {path}")

    audit = read_json(audit_path)
    expected = {
        "status": "ok",
        "method": method,
        "seed": seed,
        "split_role": "final_test",
        "annotation_sha256": EXPECTED_TEST_ANNOTATION_SHA256,
    }
    for key, value in expected.items():
        if audit.get(key) != value:
            raise ValueError(f"Protocol audit mismatch {key}: {audit.get(key)!r} != {value!r}")
    max_dets = {int(value) for value in audit.get("max_dets", [])}
    if 3000 not in max_dets:
        raise ValueError(f"Protocol audit lacks maxDets=3000: {audit_path}")
    prediction_hashes = audit.get("prediction_sha256", {})
    for kind, path in (("bbox", bbox_path), ("segm", segm_path)):
        recorded = str(prediction_hashes.get(kind, ""))
        if len(recorded) != 64 or any(char not in "0123456789abcdef" for char in recorded):
            raise ValueError(f"Invalid recorded prediction hash for {kind}: {audit_path}")
        if require_prediction_files and recorded != sha256(path):
            raise ValueError(f"Prediction hash mismatch: {path}")

    rows = read_csv(csv_path)
    bbox = select_all_3000(rows, "bbox")
    segm = select_all_3000(rows, "segm")
    evidence = {
        "test_dir": str(test_dir.resolve()),
        "protocol_csv_sha256": sha256(csv_path),
        "protocol_audit_sha256": sha256(audit_path),
        "bbox_predictions_sha256": prediction_hashes["bbox"],
        "segm_predictions_sha256": prediction_hashes["segm"],
        "prediction_files_rehashed": require_prediction_files,
    }
    return bbox, segm, evidence


def validate_seed44_training(run_dir: Path, method: str) -> dict:
    summary_path = run_dir / "summary.json"
    config_path = run_dir / "run_config.json"
    checkpoint_path = run_dir / "checkpoint_last.pth"
    inference_path = run_dir / "inference_complete.json"
    for path in (summary_path, config_path, checkpoint_path, inference_path):
        if not path.is_file() or path.stat().st_size == 0:
            raise FileNotFoundError(f"Missing seed-44 training artifact: {path}")

    summary = read_json(summary_path)
    config = read_json(config_path)
    inference = read_json(inference_path)
    args = config.get("args", {})
    protocol = config.get("protocol", {})
    expected_args = {
        "epochs": 20,
        "seed": 44,
        "batch_size": 1,
        "num_workers": 4,
        "min_size": 512,
        "max_size": 768,
        "train_crop_size": 0,
        "max_train_images": 0,
        "max_val_images": 0,
        "eval_max_images": 0,
        "eval_max_predictions": 3000,
        "box_detections_per_image": 3000,
        "rpn_pre_nms_top_n_test": 6000,
        "rpn_post_nms_top_n_test": 3000,
        "livecell_dense_protocol": True,
        "defer_coco_eval": True,
    }
    for key, expected in expected_args.items():
        if str(args.get(key)) != str(expected):
            raise ValueError(
                f"Seed-44 config mismatch for {method}: {key}={args.get(key)!r}, expected={expected!r}"
            )
    for key, expected in (
        ("lr", 1e-4),
        ("eval_score_threshold", 0.05),
        ("backbone_lr_multiplier", METHODS[method]["multiplier"]),
    ):
        actual = float(args.get(key, float("nan")))
        if not math.isclose(actual, expected, rel_tol=0.0, abs_tol=1e-12):
            raise ValueError(f"Seed-44 numeric config mismatch: {method} {key}={actual}")

    if protocol.get("name") != "livecell_dense_eval_compatible_v1":
        raise ValueError(f"Wrong dense protocol for {method}: {protocol}")
    for split in ("train_dataset", "val_dataset"):
        if config.get(split, {}).get("instance_limit") != "all":
            raise ValueError(f"Instance cap detected in {method} {split}")
    if int(config["train_dataset"]["num_images"]) != 162:
        raise ValueError(f"Unexpected seed-44 training image count for {method}")
    if int(config["val_dataset"]["num_images"]) != 570:
        raise ValueError(f"Unexpected seed-44 validation image count for {method}")

    if method == "coco_full":
        if not args.get("pretrained_coco") or any(
            args.get(flag) for flag in ("imagenet_coco_heads", "imagenet_backbone")
        ):
            raise ValueError(f"Wrong COCO-full initialization flags: {args}")
    elif method == "imagenet_coco":
        if not args.get("imagenet_coco_heads") or any(
            args.get(flag) for flag in ("pretrained_coco", "imagenet_backbone")
        ):
            raise ValueError(f"Wrong ImageNet+COCO initialization flags: {args}")
    else:
        if not args.get("hybrid_coco_heads") or not config.get("hybrid_coco_heads"):
            raise ValueError(f"Wrong SSL hybrid initialization flags for {method}: {args}")
        backbone_path = Path(str(args.get("ssl_backbone", "")))
        expected_parent = {
            "simclr_coco": "ssl_simclr_v2",
            "supcon_coco": "ssl_supcon_v2",
        }[method]
        if (
            backbone_path.name != "backbone_resnet50_ssl.pth"
            or backbone_path.parent.name != expected_parent
        ):
            raise ValueError(f"Unexpected SSL backbone path for {method}: {backbone_path}")

    checkpoint_hash = sha256(checkpoint_path)
    if summary.get("status") != "ok" or summary.get("checkpoint_sha256") != checkpoint_hash:
        raise ValueError(f"Seed-44 summary/checkpoint mismatch for {method}")
    if inference.get("status") != "ok" or inference.get("metrics_status") != "deferred":
        raise ValueError(f"Seed-44 validation inference marker is invalid for {method}")
    if inference.get("checkpoint_sha256") != checkpoint_hash:
        raise ValueError(f"Seed-44 validation marker/checkpoint mismatch for {method}")
    return {
        "run_config_sha256": sha256(config_path),
        "summary_sha256": sha256(summary_path),
        "checkpoint_sha256": checkpoint_hash,
        "validation_inference_sha256": sha256(inference_path),
    }


def write_csv(path: Path, rows: list[dict], fields: list[str]) -> None:
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    args = parse_args()
    for path in (args.old_eval_root, args.seed44_root):
        if not path.is_dir():
            raise FileNotFoundError(path)
    args.output_dir.mkdir(parents=True, exist_ok=True)

    raw_rows: list[dict] = []
    source_evidence: dict[str, dict] = {}
    for seed in (42, 43, 44):
        for method, spec in METHODS.items():
            if seed < 44:
                test_dir = args.old_eval_root / f"test_{method}_s{seed}"
                training_evidence = None
            else:
                run_dir = args.seed44_root / f"{method}_e20_s44_bblr{spec['multiplier']}"
                training_evidence = validate_seed44_training(run_dir, method)
                test_dir = run_dir / "test_official"
            bbox, segm, protocol_evidence = validate_protocol_evidence(
                test_dir=test_dir,
                method=method,
                seed=seed,
                require_prediction_files=(seed == 44),
            )
            if seed == 44:
                test_inference_path = test_dir / "inference_complete.json"
                test_inference = read_json(test_inference_path)
                checkpoint_hash = training_evidence["checkpoint_sha256"]
                if (
                    test_inference.get("status") != "ok"
                    or test_inference.get("split_role") != "final_test"
                    or test_inference.get("annotation_sha256")
                    != EXPECTED_TEST_ANNOTATION_SHA256
                    or test_inference.get("checkpoint_sha256") != checkpoint_hash
                ):
                    raise ValueError(f"Invalid seed-44 official-test inference audit: {test_inference_path}")
                protocol_evidence["test_inference_sha256"] = sha256(test_inference_path)
                protocol_evidence["training"] = training_evidence

            raw_rows.append(
                {
                    "method": spec["label"],
                    "method_key": method,
                    "backbone_lr_multiplier": spec["multiplier"],
                    "seed": seed,
                    "segm_AP": segm["AP"],
                    "segm_AP50": segm["AP50"],
                    "segm_AR": segm["AR"],
                    "bbox_AP": bbox["AP"],
                    "bbox_AP50": bbox["AP50"],
                    "bbox_AR": bbox["AR"],
                }
            )
            source_evidence[f"{method}/seed{seed}"] = protocol_evidence

    raw_rows.sort(key=lambda row: (list(METHODS).index(row["method_key"]), row["seed"]))
    raw_path = args.output_dir / "official_test_raw_n3.csv"
    raw_fields = [
        "method", "method_key", "backbone_lr_multiplier", "seed",
        "segm_AP", "segm_AP50", "segm_AR", "bbox_AP", "bbox_AP50", "bbox_AR",
    ]
    write_csv(raw_path, raw_rows, raw_fields)

    summary_rows: list[dict] = []
    for method, spec in METHODS.items():
        subset = [row for row in raw_rows if row["method_key"] == method]
        if [row["seed"] for row in subset] != [42, 43, 44]:
            raise ValueError(f"Incomplete three-seed coverage for {method}")
        summary = {
            "method": spec["label"],
            "method_key": method,
            "backbone_lr_multiplier": spec["multiplier"],
            "n_seeds": 3,
        }
        for metric in ("segm_AP", "segm_AP50", "segm_AR", "bbox_AP"):
            mean, sd = mean_sample_sd([float(row[metric]) for row in subset])
            summary[f"{metric}_mean"] = mean
            summary[f"{metric}_sample_SD"] = sd
        summary_rows.append(summary)
    summary_path = args.output_dir / "official_test_summary_n3.csv"
    write_csv(summary_path, summary_rows, list(summary_rows[0]))

    index = {
        (row["method_key"], int(row["seed"])): float(row["segm_AP"])
        for row in raw_rows
    }
    comparisons = (
        ("coco_full", "imagenet_coco"),
        ("coco_full", "simclr_coco"),
        ("coco_full", "supcon_coco"),
        ("simclr_coco", "imagenet_coco"),
        ("simclr_coco", "supcon_coco"),
    )
    difference_rows: list[dict] = []
    for left, right in comparisons:
        differences = [index[(left, seed)] - index[(right, seed)] for seed in (42, 43, 44)]
        mean, sd = mean_sample_sd(differences)
        difference_rows.append(
            {
                "comparison": f"{METHODS[left]['label']}-{METHODS[right]['label']}",
                "seed42_delta_AP": differences[0],
                "seed43_delta_AP": differences[1],
                "seed44_delta_AP": differences[2],
                "mean_delta_AP": mean,
                "sample_SD_delta_AP": sd,
                "same_direction_all_three": all(value > 0 for value in differences)
                or all(value < 0 for value in differences),
            }
        )
    differences_path = args.output_dir / "official_test_paired_differences_n3.csv"
    write_csv(differences_path, difference_rows, list(difference_rows[0]))

    generated = (raw_path, summary_path, differences_path)
    audit = {
        "status": "ok",
        "scope": "official LIVECell test, maxDets=3000, seeds 42/43/44",
        "expected_test_annotation_sha256": EXPECTED_TEST_ANNOTATION_SHA256,
        "runs": len(raw_rows),
        "methods": list(METHODS),
        "seeds": [42, 43, 44],
        "source_evidence": source_evidence,
        "generated_sha256": {path.name: sha256(path) for path in generated},
    }
    audit_path = args.output_dir / "submission_seed44_collection_audit.json"
    audit_path.write_text(json.dumps(audit, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(audit, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
