#!/usr/bin/env python3
"""Train and evaluate a LIVECell Mask R-CNN baseline with COCO metrics."""

from __future__ import annotations

import argparse
import csv
import json
import random
import time
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np
import torch
from PIL import Image
from pycocotools import mask as mask_utils
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from torch.utils.data import DataLoader, Dataset

from adaptation_optimizer import (
    build_adaptation_optimizer,
    optimizer_group_summary,
    sha256_file,
    snapshot_trainable_parameters,
    validate_multiplier,
    write_parameter_adaptation,
)
from livecell_protocol import (
    LIVECELL_AREA_RANGE_LABELS,
    LIVECELL_AREA_RANGES,
    apply_livecell_dense_protocol,
    dense_protocol_manifest,
)

from train_maskrcnn_smoke import (
    annotation_to_mask,
    build_model,
    collate_fn,
    draw_masks,
    find_image_path,
    index_images,
    move_targets_to_device,
    pil_to_tensor,
    set_seed,
    write_csv,
)


SIZE_SPLIT_NAMES = {
    "2": "0_train2percent.json",
    "4": "1_train4percent.json",
    "5": "2_train5percent.json",
    "25": "3_train25percent.json",
    "50": "4_train50percent.json",
}

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-root", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--train-ann-file", type=Path)
    parser.add_argument("--train-size-split", choices=sorted(SIZE_SPLIT_NAMES), default="5")
    parser.add_argument("--val-ann-file", type=Path)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--max-steps", type=int, default=0, help="0 means no explicit step limit.")
    parser.add_argument("--max-train-images", type=int, default=0, help="Dry-run limit. 0 uses all selected train images.")
    parser.add_argument("--max-val-images", type=int, default=0, help="Dry-run limit. 0 uses all selected val images.")
    parser.add_argument("--max-instances-per-image", type=int, default=128)
    parser.add_argument(
        "--livecell-dense-protocol",
        action="store_true",
        help=(
            "Use all annotations, a dense training crop, and maxDets=3000. "
            "Legacy defaults remain unchanged when this flag is absent."
        ),
    )
    parser.add_argument(
        "--train-crop-size",
        type=int,
        default=-1,
        help="Training-only square crop; -1 selects dense default, 0 disables cropping.",
    )
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--lr", type=float, default=0.0001)
    parser.add_argument("--min-size", type=int, default=512)
    parser.add_argument("--max-size", type=int, default=768)
    parser.add_argument("--pretrained-coco", action="store_true")
    parser.add_argument("--imagenet-backbone", action="store_true",
                        help="Use ImageNet-pretrained backbone + random FPN/RPN/heads (fair baseline vs SSL)")
    parser.add_argument("--imagenet-coco-heads", action="store_true",
                        help="Use ImageNet V2 backbone + COCO-pretrained FPN/RPN/ROI heads")
    parser.add_argument(
        "--coco-backbone-random-components",
        action="store_true",
        help=(
            "Use only the COCO-adapted ResNet backbone.body and randomly initialize "
            "FPN/RPN/RoI components. This isolates COCO backbone adaptation from "
            "jointly pretrained detector components."
        ),
    )
    parser.add_argument("--freeze-backbone", action="store_true",
                        help="Freeze backbone.body, train only FPN/RPN/ROI heads")
    parser.add_argument("--freeze-heads", action="store_true",
                        help=(
                            "Freeze FPN/RPN/ROI intermediate heads, train backbone.body "
                            "+ final LIVECell predictors"
                        ))
    parser.add_argument("--eval-score-threshold", type=float, default=0.05)
    parser.add_argument("--eval-max-predictions", type=int, default=100)
    parser.add_argument("--box-detections-per-image", type=int, default=100)
    parser.add_argument("--rpn-pre-nms-top-n-train", type=int, default=1000)
    parser.add_argument("--rpn-post-nms-top-n-train", type=int, default=500)
    parser.add_argument("--rpn-pre-nms-top-n-test", type=int, default=500)
    parser.add_argument("--rpn-post-nms-top-n-test", type=int, default=200)
    parser.add_argument("--box-batch-size-per-image", type=int, default=128)
    parser.add_argument("--eval-max-images", type=int, default=0, help="0 evaluates all validation images.")
    parser.add_argument("--vis-count", type=int, default=8)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"])
    parser.add_argument("--defer-coco-eval", action="store_true",
                        help="Skip built-in COCOeval; write predictions only for external protocol ablation.")
    parser.add_argument(
        "--backbone-lr-multiplier",
        type=validate_multiplier,
        default=1.0,
        help=(
            "Multiplier applied to --lr for backbone.body.* parameters only. "
            "FPN/RPN/ROI heads keep the base LR. Default 1.0 (uniform LR)."
        ),
    )
    return parser.parse_args()


def discover_size_split(data_root: Path, percent: str) -> Path:
    expected = SIZE_SPLIT_NAMES[percent]
    candidates = sorted(data_root.rglob(expected))
    if candidates:
        return candidates[0]
    candidates = sorted(data_root.rglob(f"*train{percent}percent*.json"))
    if candidates:
        return candidates[0]
    raise FileNotFoundError(f"Could not find LIVECell size split for {percent}% under {data_root}")


def discover_val_ann_file(data_root: Path) -> Path:
    candidates = [
        p
        for p in sorted(data_root.rglob("*.json"))
        if "val" in p.name.lower()
        and "coco" in p.name.lower()
        and "size_split" not in str(p).lower()
        and "single_cells" not in str(p).lower()
    ]
    if candidates:
        return candidates[0]
    candidates = [p for p in sorted(data_root.rglob("*.json")) if "val" in p.name.lower()]
    if candidates:
        return candidates[0]
    raise FileNotFoundError(f"Could not find validation COCO annotation under {data_root}")


def load_json(path: Path) -> Dict[str, Any]:
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def parameter_group_name(name: str) -> str:
    if name.startswith("backbone.body."):
        return "backbone.body"
    if name.startswith("backbone.fpn."):
        return "backbone.fpn"
    if name.startswith("rpn."):
        return "rpn"
    if name.startswith("roi_heads.box_predictor."):
        return "roi_heads.box_predictor"
    if name.startswith("roi_heads.mask_predictor."):
        return "roi_heads.mask_predictor"
    if name.startswith("roi_heads."):
        return "roi_heads.other"
    return name.split(".", 1)[0]


def summarize_parameter_training(model: torch.nn.Module) -> Dict[str, Any]:
    summary: Dict[str, Any] = {
        "total_tensors": 0,
        "trainable_tensors": 0,
        "frozen_tensors": 0,
        "total_params": 0,
        "trainable_params": 0,
        "frozen_params": 0,
        "groups": {},
    }
    groups: Dict[str, Dict[str, int]] = defaultdict(
        lambda: {
            "total_tensors": 0,
            "trainable_tensors": 0,
            "frozen_tensors": 0,
            "total_params": 0,
            "trainable_params": 0,
            "frozen_params": 0,
        }
    )
    for name, param in model.named_parameters():
        group = groups[parameter_group_name(name)]
        n_params = int(param.numel())
        summary["total_tensors"] += 1
        summary["total_params"] += n_params
        group["total_tensors"] += 1
        group["total_params"] += n_params
        if param.requires_grad:
            summary["trainable_tensors"] += 1
            summary["trainable_params"] += n_params
            group["trainable_tensors"] += 1
            group["trainable_params"] += n_params
        else:
            summary["frozen_tensors"] += 1
            summary["frozen_params"] += n_params
            group["frozen_tensors"] += 1
            group["frozen_params"] += n_params
    summary["groups"] = dict(sorted(groups.items()))
    return summary


def build_coco_backbone_random_components_model(args: argparse.Namespace) -> torch.nn.Module:
    """Build matched random detector components with a COCO-trained body.

    The target starts from the ImageNet-backbone/random-component factory path so
    its FrozenBatchNorm and trainable-layer policy match the other pretrained
    backbone controls. Building a fully random detector here would silently use
    BatchNorm2d and make all five ResNet stages trainable in torchvision.
    """
    common = {
        "num_classes": 2,
        "min_size": args.min_size,
        "max_size": args.max_size,
        "box_detections_per_img": args.box_detections_per_image,
        "rpn_pre_nms_top_n_train": args.rpn_pre_nms_top_n_train,
        "rpn_post_nms_top_n_train": args.rpn_post_nms_top_n_train,
        "rpn_pre_nms_top_n_test": args.rpn_pre_nms_top_n_test,
        "rpn_post_nms_top_n_test": args.rpn_post_nms_top_n_test,
        "box_batch_size_per_image": args.box_batch_size_per_image,
    }
    target = build_model(**common, imagenet_backbone=True)
    source = build_model(**common, pretrained_coco=True)
    mapped = {
        name: tensor.detach().cpu()
        for name, tensor in source.state_dict().items()
        if name.startswith("backbone.body.")
    }
    if not mapped:
        raise RuntimeError("No backbone.body.* keys found in COCO source model")
    missing, unexpected = target.load_state_dict(mapped, strict=False)
    backbone_missing = [name for name in missing if name.startswith("backbone.body.")]
    if backbone_missing or unexpected:
        raise RuntimeError(
            "COCO-backbone transfer failed: "
            f"backbone_missing={backbone_missing[:8]}, unexpected={unexpected[:8]}"
        )
    body_training = summarize_parameter_training(target)["groups"]["backbone.body"]
    if not body_training["trainable_params"] or not body_training["frozen_params"]:
        raise RuntimeError(
            "Matched COCO-backbone control must retain the pretrained detector's "
            f"partial-freeze policy, got {body_training}"
        )
    del source
    print(
        "COCO_BACKBONE_RANDOM_COMPONENTS_LOADED",
        json.dumps({
            "transferred_tensors": len(mapped),
            "detector_components": "random",
            "construction_reference": "imagenet_backbone_random_components",
            "backbone_training": body_training,
        }),
    )
    return target


class LiveCellCocoDataset(Dataset):
    def __init__(
        self,
        data_root: Path,
        annotation_file: Path,
        image_index: Dict[str, List[Path]],
        max_images: int,
        max_instances_per_image: int,
        seed: int,
        training: bool = False,
        crop_size: int = 0,
    ) -> None:
        self.data_root = data_root
        self.annotation_file = annotation_file
        self.coco_dict = load_json(annotation_file)
        self.max_instances_per_image = max_instances_per_image
        self.training = training
        self.crop_size = max(0, int(crop_size))

        anns_by_image: Dict[int, List[Dict[str, Any]]] = defaultdict(list)
        for ann in self.coco_dict.get("annotations", []):
            anns_by_image[int(ann["image_id"])].append(ann)
        self.anns_by_image = anns_by_image

        items: List[Dict[str, Any]] = []
        for img in self.coco_dict.get("images", []):
            image_id = int(img["id"])
            if image_id not in anns_by_image:
                continue
            image_path = find_image_path(str(img.get("file_name", "")), image_index, data_root)
            if image_path is None:
                continue
            items.append(
                {
                    "image": img,
                    "image_path": image_path,
                    "annotation_count": len(anns_by_image[image_id]),
                }
            )

        rng = random.Random(seed)
        if max_images > 0:
            rng.shuffle(items)
            items = items[:max_images]
        self.items = items

        if not self.items:
            raise RuntimeError(f"No usable images found for annotation file: {annotation_file}")

    def __len__(self) -> int:
        return len(self.items)

    def _sample_crop(
        self,
        width: int,
        height: int,
        annotations: Sequence[Dict[str, Any]],
    ) -> Tuple[int, int, int, int]:
        crop_width = min(self.crop_size, width)
        crop_height = min(self.crop_size, height)
        if not self.training or self.crop_size <= 0 or (crop_width == width and crop_height == height):
            return 0, 0, width, height

        max_left = width - crop_width
        max_top = height - crop_height
        for _ in range(20):
            left = random.randint(0, max_left) if max_left > 0 else 0
            top = random.randint(0, max_top) if max_top > 0 else 0
            right = left + crop_width
            bottom = top + crop_height
            if any(self._annotation_intersects_crop(ann, left, top, right, bottom) for ann in annotations):
                return left, top, right, bottom

        valid = [ann for ann in annotations if ann.get("bbox") and len(ann["bbox"]) == 4]
        if valid:
            x, y, box_width, box_height = [float(value) for value in random.choice(valid)["bbox"]]
            center_x = x + box_width / 2.0
            center_y = y + box_height / 2.0
            left = min(max(int(round(center_x - crop_width / 2.0)), 0), max_left)
            top = min(max(int(round(center_y - crop_height / 2.0)), 0), max_top)
            return left, top, left + crop_width, top + crop_height
        return 0, 0, crop_width, crop_height

    @staticmethod
    def _annotation_intersects_crop(
        annotation: Dict[str, Any],
        left: int,
        top: int,
        right: int,
        bottom: int,
    ) -> bool:
        bbox = annotation.get("bbox")
        if not bbox or len(bbox) != 4:
            return False
        x, y, width, height = [float(value) for value in bbox]
        return x + width > left and y + height > top and x < right and y < bottom

    def __getitem__(self, index: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        item = self.items[index]
        img_info = item["image"]
        image_id = int(img_info["id"])
        image_path = item["image_path"]

        annotations = list(self.anns_by_image.get(image_id, []))
        if self.max_instances_per_image > 0:
            annotations.sort(key=lambda ann: float(ann.get("area", 0.0)), reverse=True)
            annotations = annotations[: self.max_instances_per_image]

        with Image.open(image_path) as image:
            image = image.convert("RGB")
            source_width, source_height = image.size
            left, top, right, bottom = self._sample_crop(source_width, source_height, annotations)
            image = image.crop((left, top, right, bottom))
            width, height = image.size
            image_tensor = pil_to_tensor(image)

        boxes: List[List[float]] = []
        masks: List[np.ndarray] = []
        areas: List[float] = []
        iscrowd: List[int] = []

        for ann in annotations:
            if not self._annotation_intersects_crop(ann, left, top, right, bottom):
                continue
            bbox = ann.get("bbox")
            if not bbox or len(bbox) != 4:
                continue
            x, y, w, h = [float(v) for v in bbox]
            mask = annotation_to_mask(ann, source_height, source_width)
            if mask.shape != (source_height, source_width):
                if mask.T.shape == (source_height, source_width):
                    mask = mask.T
                else:
                    continue
            mask = mask[top:bottom, left:right]
            if not mask.any():
                continue

            crop_applied = left != 0 or top != 0 or right != source_width or bottom != source_height
            if crop_applied:
                mask_y, mask_x = np.nonzero(mask)
                x1 = float(mask_x.min())
                y1 = float(mask_y.min())
                x2 = float(mask_x.max() + 1)
                y2 = float(mask_y.max() + 1)
                area = float(mask.sum())
            else:
                x1 = max(0.0, x)
                y1 = max(0.0, y)
                x2 = min(float(width), x + max(0.0, w))
                y2 = min(float(height), y + max(0.0, h))
                if x2 <= x1 or y2 <= y1:
                    continue
                area = float(ann.get("area", w * h))
            boxes.append([x1, y1, x2, y2])
            masks.append(mask)
            areas.append(area)
            iscrowd.append(int(ann.get("iscrowd", 0)))

        if not boxes:
            boxes = [[0.0, 0.0, 1.0, 1.0]]
            masks = [np.zeros((height, width), dtype=np.uint8)]
            areas = [1.0]
            iscrowd = [1]

        target = {
            "boxes": torch.as_tensor(boxes, dtype=torch.float32),
            "labels": torch.ones((len(boxes),), dtype=torch.int64),
            "masks": torch.as_tensor(np.stack(masks), dtype=torch.uint8),
            "image_id": torch.as_tensor([image_id], dtype=torch.int64),
            "area": torch.as_tensor(areas, dtype=torch.float32),
            "iscrowd": torch.as_tensor(iscrowd, dtype=torch.int64),
        }
        return image_tensor, target

    def describe(self) -> Dict[str, Any]:
        counts = [int(item["annotation_count"]) for item in self.items]
        return {
            "annotation_file": str(self.annotation_file),
            "num_images": len(self.items),
            "min_annotations_per_selected_image": min(counts),
            "max_annotations_per_selected_image": max(counts),
            "mean_annotations_per_selected_image": float(np.mean(counts)),
            "total_annotations_in_selected_images": int(sum(counts)),
            "max_instances_per_image_used": self.max_instances_per_image,
            "instance_limit": "all" if self.max_instances_per_image == 0 else self.max_instances_per_image,
            "training": self.training,
            "crop_size": self.crop_size,
        }


def train_model(
    model: torch.nn.Module,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    epochs: int,
    max_steps: int,
) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    global_step = 0
    model.train()

    for epoch in range(epochs):
        for images, targets in loader:
            started = time.time()
            images = [image.to(device) for image in images]
            targets = move_targets_to_device(targets, device)

            loss_dict = model(images, targets)
            loss = sum(value for value in loss_dict.values())
            if not torch.isfinite(loss):
                raise RuntimeError(f"Non-finite loss at step {global_step}: {loss.item()}")

            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()

            row = {
                "epoch": epoch,
                "step": global_step,
                "total_loss": float(loss.detach().cpu()),
                "seconds": round(time.time() - started, 3),
            }
            row.update({name: float(value.detach().cpu()) for name, value in loss_dict.items()})
            if device.type == "cuda":
                row["cuda_memory_mb"] = round(torch.cuda.max_memory_allocated(device) / 1024 / 1024, 1)
            rows.append(row)

            print(
                "TRAIN_STEP",
                global_step,
                "epoch",
                epoch,
                "loss",
                f"{row['total_loss']:.4f}",
                "time_s",
                row["seconds"],
                flush=True,
            )

            global_step += 1
            if max_steps > 0 and global_step >= max_steps:
                return rows

    return rows


def mask_to_coco_rle(mask: np.ndarray) -> Dict[str, Any]:
    rle = mask_utils.encode(np.asfortranarray(mask.astype(np.uint8)))
    rle["counts"] = rle["counts"].decode("ascii")
    return rle


def _mean_valid(values: np.ndarray) -> float:
    values = values[values > -1]
    return float(np.mean(values)) if values.size else -1.0


def _coco_metric(
    evaluator: COCOeval,
    use_precision: bool,
    iou_threshold: Optional[float] = None,
    area_label: str = "all",
    max_detections: Optional[int] = None,
) -> float:
    params = evaluator.params
    max_detections = max_detections if max_detections is not None else params.maxDets[-1]
    area_index = params.areaRngLbl.index(area_label)
    max_dets_index = params.maxDets.index(max_detections)

    if use_precision:
        values = evaluator.eval["precision"][:, :, :, area_index, max_dets_index]
    else:
        values = evaluator.eval["recall"][:, :, area_index, max_dets_index]
    if iou_threshold is not None:
        iou_indices = np.where(np.isclose(params.iouThrs, iou_threshold))[0]
        if iou_indices.size == 0:
            return -1.0
        values = values[iou_indices]
    return _mean_valid(values)


def summarize_coco_eval(evaluator: COCOeval) -> Dict[str, Any]:
    """Summarize COCO arrays without pycocotools' hard-coded maxDets=100."""
    max_detections = int(evaluator.params.maxDets[-1])
    ar_max = _coco_metric(evaluator, use_precision=False, max_detections=max_detections)
    metrics = {
        "AP": _coco_metric(evaluator, use_precision=True, max_detections=max_detections),
        "AP50": _coco_metric(
            evaluator, use_precision=True, iou_threshold=0.50, max_detections=max_detections
        ),
        "AP75": _coco_metric(
            evaluator, use_precision=True, iou_threshold=0.75, max_detections=max_detections
        ),
        "AP_small": _coco_metric(
            evaluator, use_precision=True, area_label="small", max_detections=max_detections
        ),
        "AP_medium": _coco_metric(
            evaluator, use_precision=True, area_label="medium", max_detections=max_detections
        ),
        "AP_large": _coco_metric(
            evaluator, use_precision=True, area_label="large", max_detections=max_detections
        ),
        "AR_1": _coco_metric(evaluator, use_precision=False, max_detections=1),
        "AR_10": _coco_metric(evaluator, use_precision=False, max_detections=10),
        f"AR_{max_detections}": ar_max,
        "AR_max": ar_max,
        "AR_small": _coco_metric(
            evaluator, use_precision=False, area_label="small", max_detections=max_detections
        ),
        "AR_medium": _coco_metric(
            evaluator, use_precision=False, area_label="medium", max_detections=max_detections
        ),
        "AR_large": _coco_metric(
            evaluator, use_precision=False, area_label="large", max_detections=max_detections
        ),
        "AFNR": 1.0 - ar_max if ar_max >= 0.0 else -1.0,
        "max_detections_per_image": max_detections,
    }
    return metrics


def run_coco_eval(
    annotation_file: Path,
    predictions: List[Dict[str, Any]],
    iou_type: str,
    max_detections: int = 100,
    image_ids: Optional[Sequence[int]] = None,
    use_livecell_area_ranges: bool = False,
) -> Dict[str, Any]:
    if max_detections <= 10:
        raise ValueError("max_detections must be greater than 10 for COCO evaluation")
    if not predictions:
        return {
            "status": "no_predictions",
            "num_predictions": 0,
            "max_detections_per_image": max_detections,
            "area_ranges": "LIVECell_500_1500" if use_livecell_area_ranges else "COCO_default",
            "AP": 0.0,
            "AFNR": 1.0,
        }

    coco_gt = COCO(str(annotation_file))
    coco_dt = coco_gt.loadRes(predictions)
    evaluator = COCOeval(coco_gt, coco_dt, iouType=iou_type)
    evaluator.params.maxDets = [1, 10, int(max_detections)]
    if use_livecell_area_ranges:
        evaluator.params.areaRng = LIVECELL_AREA_RANGES
        evaluator.params.areaRngLbl = LIVECELL_AREA_RANGE_LABELS
    if image_ids is not None:
        evaluator.params.imgIds = sorted({int(image_id) for image_id in image_ids})
    evaluator.evaluate()
    evaluator.accumulate()
    metrics = summarize_coco_eval(evaluator)
    metrics["status"] = "ok"
    metrics["num_predictions"] = len(predictions)
    metrics["area_ranges"] = "LIVECell_500_1500" if use_livecell_area_ranges else "COCO_default"
    return metrics


@torch.no_grad()
def evaluate_and_visualize(
    model: torch.nn.Module,
    dataset: LiveCellCocoDataset,
    device: torch.device,
    output_dir: Path,
    score_threshold: float,
    max_predictions: int,
    eval_max_images: int,
    vis_count: int,
    use_livecell_area_ranges: bool = False,
    defer_coco_eval: bool = False,
) -> Dict[str, Any]:
    model.eval()
    bbox_predictions: List[Dict[str, Any]] = []
    segm_predictions: List[Dict[str, Any]] = []
    quick_rows: List[Dict[str, Any]] = []
    evaluated_image_ids: List[int] = []

    limit = len(dataset) if eval_max_images <= 0 else min(eval_max_images, len(dataset))

    for index in range(limit):
        image, target = dataset[index]
        image_id = int(target["image_id"].item())
        evaluated_image_ids.append(image_id)
        prediction = model([image.to(device)])[0]
        scores = prediction["scores"].detach().cpu()
        keep = scores >= score_threshold

        pred_scores = scores[keep].numpy()
        pred_boxes_xyxy = prediction["boxes"].detach().cpu()[keep].numpy()
        pred_masks = (prediction["masks"].detach().cpu()[keep, 0].numpy() >= 0.5).astype(np.uint8)

        if len(pred_scores) > max_predictions:
            order = np.argsort(-pred_scores)[:max_predictions]
            pred_scores = pred_scores[order]
            pred_boxes_xyxy = pred_boxes_xyxy[order]
            pred_masks = pred_masks[order]

        for score, box_xyxy, mask in zip(pred_scores, pred_boxes_xyxy, pred_masks):
            x1, y1, x2, y2 = [float(v) for v in box_xyxy]
            bbox_xywh = [x1, y1, max(0.0, x2 - x1), max(0.0, y2 - y1)]
            record = {
                "image_id": image_id,
                "category_id": 1,
                "bbox": bbox_xywh,
                "score": float(score),
            }
            bbox_predictions.append(record)
            segm_record = dict(record)
            segm_record["segmentation"] = mask_to_coco_rle(mask)
            segm_predictions.append(segm_record)

        quick_rows.append(
            {
                "val_index": index,
                "image_id": image_id,
                "gt_instances": int(target["masks"].shape[0]),
                "predictions_above_threshold": int(len(pred_scores)),
                "top_score": float(pred_scores[0]) if len(pred_scores) else 0.0,
            }
        )

        if index < vis_count:
            image_array = (image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
            draw_masks(
                image_array,
                target["masks"].cpu().numpy(),
                target["boxes"].cpu().numpy(),
                None,
                output_dir / "visualization" / f"val_{index:03d}_gt.png",
                max_items=max_predictions,
            )
            draw_masks(
                image_array,
                pred_masks,
                pred_boxes_xyxy,
                pred_scores,
                output_dir / "visualization" / f"val_{index:03d}_pred.png",
                max_items=max_predictions,
            )

        print("EVAL_IMAGE", index, "pred", len(pred_scores), flush=True)

    (output_dir / "predictions_bbox.json").write_text(json.dumps(bbox_predictions), encoding="utf-8")
    (output_dir / "predictions_segm.json").write_text(json.dumps(segm_predictions), encoding="utf-8")
    write_csv(output_dir / "quick_eval.csv", quick_rows)

    if defer_coco_eval:
        return {
            "metrics_status": "deferred",
            "num_eval_images": limit,
            "num_bbox_predictions": len(bbox_predictions),
            "num_segm_predictions": len(segm_predictions),
        }

    bbox_metrics = run_coco_eval(
        dataset.annotation_file,
        bbox_predictions,
        "bbox",
        max_detections=max_predictions,
        image_ids=evaluated_image_ids,
        use_livecell_area_ranges=use_livecell_area_ranges,
    )
    segm_metrics = run_coco_eval(
        dataset.annotation_file,
        segm_predictions,
        "segm",
        max_detections=max_predictions,
        image_ids=evaluated_image_ids,
        use_livecell_area_ranges=use_livecell_area_ranges,
    )
    (output_dir / "coco_eval_bbox.json").write_text(json.dumps(bbox_metrics, indent=2), encoding="utf-8")
    (output_dir / "coco_eval_segm.json").write_text(json.dumps(segm_metrics, indent=2), encoding="utf-8")

    return {
        "bbox": bbox_metrics,
        "segm": segm_metrics,
        "num_eval_images": limit,
        "num_bbox_predictions": len(bbox_predictions),
        "num_segm_predictions": len(segm_predictions),
    }


def main() -> None:
    args = apply_livecell_dense_protocol(parse_args())
    set_seed(args.seed)
    args.output_dir.mkdir(parents=True, exist_ok=True)

    if args.device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)

    train_ann_file = args.train_ann_file or discover_size_split(args.data_root, args.train_size_split)
    val_ann_file = args.val_ann_file or discover_val_ann_file(args.data_root)

    print("DEVICE", device)
    if device.type == "cuda":
        print("GPU", torch.cuda.get_device_name(device))
    print("TRAIN_ANN", train_ann_file)
    print("VAL_ANN", val_ann_file)

    image_index = index_images(args.data_root)
    train_dataset = LiveCellCocoDataset(
        args.data_root,
        train_ann_file,
        image_index,
        args.max_train_images,
        args.max_instances_per_image,
        args.seed,
        training=True,
        crop_size=args.train_crop_size,
    )
    val_dataset = LiveCellCocoDataset(
        args.data_root,
        val_ann_file,
        image_index,
        args.max_val_images,
        0 if args.livecell_dense_protocol else args.max_instances_per_image,
        args.seed + 1,
        training=False,
        crop_size=0,
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        collate_fn=collate_fn,
        pin_memory=device.type == "cuda",
    )

    initialization_flags = (
        args.pretrained_coco,
        args.imagenet_backbone,
        args.imagenet_coco_heads,
        args.coco_backbone_random_components,
    )
    if sum(bool(flag) for flag in initialization_flags) > 1:
        raise ValueError(
            "--pretrained-coco, --imagenet-backbone, --imagenet-coco-heads and "
            "--coco-backbone-random-components are mutually exclusive"
        )
    if args.coco_backbone_random_components:
        model = build_coco_backbone_random_components_model(args)
    else:
        model = build_model(
            num_classes=2,
            min_size=args.min_size,
            max_size=args.max_size,
            pretrained_coco=args.pretrained_coco,
            imagenet_backbone=args.imagenet_backbone,
            imagenet_backbone_coco_heads=args.imagenet_coco_heads,
            box_detections_per_img=args.box_detections_per_image,
            rpn_pre_nms_top_n_train=args.rpn_pre_nms_top_n_train,
            rpn_post_nms_top_n_train=args.rpn_post_nms_top_n_train,
            rpn_pre_nms_top_n_test=args.rpn_pre_nms_top_n_test,
            rpn_post_nms_top_n_test=args.rpn_post_nms_top_n_test,
            box_batch_size_per_image=args.box_batch_size_per_image,
        )
    model.to(device)

    # --- Freeze logic (mutually exclusive, only with --imagenet-coco-heads) ---
    if args.freeze_backbone or args.freeze_heads:
        if not args.imagenet_coco_heads:
            raise ValueError("--freeze-backbone/--freeze-heads requires --imagenet-coco-heads")
        if args.freeze_backbone and args.freeze_heads:
            raise ValueError("--freeze-backbone and --freeze-heads are mutually exclusive")
    if args.freeze_backbone and args.backbone_lr_multiplier != 1.0:
        raise ValueError(
            "--backbone-lr-multiplier must remain 1.0 when --freeze-backbone is used"
        )
    if args.freeze_backbone:
        for name, param in model.named_parameters():
            if name.startswith("backbone.body."):
                param.requires_grad = False
        training_summary = summarize_parameter_training(model)
        print(
            "FROZEN backbone: "
            f"{training_summary['frozen_params']} scalar params frozen, "
            f"{training_summary['trainable_params']} trainable"
        )
    if args.freeze_heads:
        for name, param in model.named_parameters():
            if name.startswith("backbone.body."):
                continue
            # The final predictors are newly created for LIVECell's 2 classes,
            # so they must remain trainable even in the freeze-heads ablation.
            if name.startswith("roi_heads.box_predictor.") or name.startswith("roi_heads.mask_predictor."):
                continue
            param.requires_grad = False
        training_summary = summarize_parameter_training(model)
        print(
            "FROZEN heads/intermediate detection layers: "
            f"{training_summary['frozen_params']} scalar params frozen, "
            f"{training_summary['trainable_params']} trainable"
        )
    training_summary = summarize_parameter_training(model)
    print("TRAINING_PARAMETER_SUMMARY", json.dumps(training_summary, ensure_ascii=False))

    optimizer = build_adaptation_optimizer(
        model,
        base_lr=args.lr,
        multiplier=args.backbone_lr_multiplier,
        momentum=0.9,
        weight_decay=0.0005,
        require_both_groups=not args.freeze_backbone,
    )
    opt_summary = optimizer_group_summary(optimizer)
    print("OPTIMIZER_GROUPS", json.dumps(opt_summary, ensure_ascii=False))

    # Snapshot initial parameters for adaptation tracking
    initial_params = snapshot_trainable_parameters(model)

    config = {
        "args": {key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()},
        "device": str(device),
        "torch": torch.__version__,
        "protocol": dense_protocol_manifest(args),
        "train_dataset": train_dataset.describe(),
        "val_dataset": val_dataset.describe(),
        "parameter_training": training_summary,
        "optimizer_groups": opt_summary,
        "backbone_lr_multiplier": args.backbone_lr_multiplier,
    }
    (args.output_dir / "run_config.json").write_text(json.dumps(config, indent=2), encoding="utf-8")
    print("TRAIN_DATASET", json.dumps(train_dataset.describe(), ensure_ascii=False))
    print("VAL_DATASET", json.dumps(val_dataset.describe(), ensure_ascii=False))

    loss_rows = train_model(model, train_loader, optimizer, device, args.epochs, args.max_steps)
    write_csv(args.output_dir / "loss_log.csv", loss_rows)

    # Compute parameter adaptation statistics
    adaptation_path = write_parameter_adaptation(args.output_dir, initial_params, model, optimizer)
    print("PARAMETER_ADAPTATION", adaptation_path)

    checkpoint_path = args.output_dir / "checkpoint_last.pth"
    torch.save(
        {
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict(),
            "config": config,
            "loss_rows": loss_rows,
        },
        checkpoint_path,
    )

    eval_summary = evaluate_and_visualize(
        model,
        val_dataset,
        device,
        args.output_dir,
        args.eval_score_threshold,
        args.eval_max_predictions,
        args.eval_max_images,
        args.vis_count,
        use_livecell_area_ranges=args.livecell_dense_protocol,
        defer_coco_eval=args.defer_coco_eval,
    )

    summary = {
        "status": "ok",
        "train_steps": len(loss_rows),
        "last_total_loss": loss_rows[-1]["total_loss"] if loss_rows else None,
        "checkpoint": str(checkpoint_path),
        "checkpoint_sha256": sha256_file(checkpoint_path),
        "parameter_adaptation": str(adaptation_path),
        "parameter_adaptation_sha256": sha256_file(adaptation_path),
        "run_config": str(args.output_dir / "run_config.json"),
        "run_config_sha256": sha256_file(args.output_dir / "run_config.json"),
        "backbone_lr_multiplier": args.backbone_lr_multiplier,
        "eval_summary": eval_summary,
    }
    (args.output_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print("BASELINE_SUMMARY", json.dumps(summary, ensure_ascii=False))

    if args.defer_coco_eval:
        import hashlib
        def _sha256(p): h=hashlib.sha256(); h.update(p.read_bytes()); return h.hexdigest()
        inference_complete = {
            "status": "ok",
            "metrics_status": "deferred",
            "predictions_bbox_sha256": _sha256(args.output_dir / "predictions_bbox.json"),
            "predictions_segm_sha256": _sha256(args.output_dir / "predictions_segm.json"),
            "checkpoint_sha256": sha256_file(checkpoint_path),
        }
        (args.output_dir / "inference_complete.json").write_text(json.dumps(inference_complete, indent=2))


if __name__ == "__main__":
    main()
