#!/usr/bin/env python3
"""Transfer SSL-pretrained ResNet50 backbone to Mask R-CNN.

Loads an SSL checkpoint, extracts the backbone weights, and builds a
Mask R-CNN model suitable for fine-tuning on LIVECell.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import torch
from torch import nn
from torchvision.models import ResNet50_Weights
from torchvision.models.detection import (
    MaskRCNN_ResNet50_FPN_Weights,
    maskrcnn_resnet50_fpn,
)
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
from torchvision.ops.misc import FrozenBatchNorm2d


def _assert_matched_pretrained_backbone_policy(model: torch.nn.Module) -> None:
    """Reject hidden norm/freeze mismatches in backbone swap experiments."""
    body = model.backbone.body
    if any(isinstance(module, nn.BatchNorm2d) for module in body.modules()):
        raise RuntimeError("Backbone swap control unexpectedly contains trainable BatchNorm2d")
    if not any(isinstance(module, FrozenBatchNorm2d) for module in body.modules()):
        raise RuntimeError("Backbone swap control is missing FrozenBatchNorm2d")
    parameters = list(body.parameters())
    if not any(parameter.requires_grad for parameter in parameters):
        raise RuntimeError("Backbone swap control has no trainable body parameters")
    if not any(not parameter.requires_grad for parameter in parameters):
        raise RuntimeError("Backbone swap control has no frozen body parameters")


def build_maskrcnn_with_ssl_backbone(
    ssl_backbone_path: str | Path,
    num_classes: int = 2,
    min_size: int = 512,
    max_size: int = 768,
    device: str = "cpu",
    box_detections_per_img: int = 100,
    rpn_pre_nms_top_n_train: int = 1000,
    rpn_post_nms_top_n_train: int = 500,
    rpn_pre_nms_top_n_test: int = 500,
    rpn_post_nms_top_n_test: int = 200,
    box_batch_size_per_image: int = 128,
) -> torch.nn.Module:
    """Build Mask R-CNN with SSL-pretrained backbone weights.

    Strategy:
    1. Create Mask R-CNN with an ImageNet backbone and random detector components
       to match FrozenBatchNorm and the three-stage trainable-layer policy used by
       the COCO-component conditions
    2. Load SSL-pretrained ResNet50 state dict
    3. Map SSL keys to Mask R-CNN backbone.body keys
    4. Replace box/mask predictors for LIVECell (1 class + background)
    """
    model = maskrcnn_resnet50_fpn(
        weights=None,
        weights_backbone=ResNet50_Weights.IMAGENET1K_V2,
        min_size=min_size,
        max_size=max_size,
        box_detections_per_img=box_detections_per_img,
        rpn_pre_nms_top_n_train=rpn_pre_nms_top_n_train,
        rpn_post_nms_top_n_train=rpn_post_nms_top_n_train,
        rpn_pre_nms_top_n_test=rpn_pre_nms_top_n_test,
        rpn_post_nms_top_n_test=rpn_post_nms_top_n_test,
        box_batch_size_per_image=box_batch_size_per_image,
    )
    _assert_matched_pretrained_backbone_policy(model)

    ssl_state = torch.load(str(ssl_backbone_path), map_location=device, weights_only=True)
    if "model" in ssl_state:
        ssl_state = ssl_state["model"]

    # Guard against passing the full ssl_checkpoint.pth instead of the
    # extracted backbone_resnet50_ssl.pth. Full checkpoints have keys like
    # "backbone.conv1.weight" which would map to "backbone.body.backbone..."
    # and silently fail to load.
    sample_keys = list(ssl_state.keys())[:3]
    if any(k.startswith("backbone.") for k in ssl_state.keys()):
        raise ValueError(
            f"Checkpoint '{ssl_backbone_path}' contains 'backbone.'-prefixed keys "
            f"(sample: {sample_keys}). This looks like the full ssl_checkpoint.pth, "
            f"not the extracted backbone. Use backbone_resnet50_ssl.pth instead."
        )

    # Filter and map keys: SSL uses standalone ResNet50,
    # Mask R-CNN expects backbone.body.{layer}
    # SSL keys: conv1.weight, bn1.weight, layer1.0.conv1.weight, etc.
    # M-RCNN keys: backbone.body.conv1.weight, backbone.body.bn1.weight,
    #               backbone.body.layer1.0.conv1.weight, etc.

    mapped = {}
    skipped = []
    for k, v in ssl_state.items():
        # Skip non-backbone keys (projection head, classifier)
        if k.startswith("encoder.") or k.startswith("fc.") or k.startswith("projection."):
            skipped.append(k)
            continue
        # Skip num_batches_tracked: SSL uses BatchNorm2d (has it),
        # Mask R-CNN uses FrozenBatchNorm2d (does NOT have it).
        # Loading it would cause unexpected key errors.
        if "num_batches_tracked" in k:
            continue
        # Map: conv1.weight -> backbone.body.conv1.weight
        mapped[f"backbone.body.{k}"] = v
    if skipped:
        print(f"Skipped {len(skipped)} non-backbone keys: {skipped[:3]}...")
    if not mapped:
        raise ValueError(f"No backbone weights were mapped from checkpoint: {ssl_backbone_path}")

    # Load mapped weights (strict=False to allow missing FPN keys)
    missing, unexpected = model.load_state_dict(mapped, strict=False)
    print(f"Loaded {len(mapped)} backbone keys from SSL checkpoint")

    # CRITICAL: verify backbone.body.* keys were actually loaded.
    # If key mapping is wrong, backbone stays randomly initialized.
    # Exclude fc.* (intentionally skipped) from the check.
    backbone_missing = [
        k for k in missing
        if k.startswith("backbone.body.") and not k.startswith("backbone.body.fc.")
    ]
    if backbone_missing:
        raise RuntimeError(
            f"{len(backbone_missing)} backbone.body.* keys are MISSING after load! "
            f"Backbone is randomly initialized. Examples: {backbone_missing[:5]}"
        )
    print(f"Missing keys (expected, FPN/RPN/heads + fc): {len(missing)}")

    if unexpected:
        raise RuntimeError(f"Unexpected keys while loading SSL backbone: {unexpected[:5]}")

    # Replace predictors for LIVECell
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = model.roi_heads.mask_predictor.conv5_mask.out_channels
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)

    return model


def build_maskrcnn_with_coco_heads_ssl_backbone(
    ssl_backbone_path: str | Path,
    num_classes: int = 2,
    min_size: int = 512,
    max_size: int = 768,
    device: str = "cpu",
    box_detections_per_img: int = 100,
    rpn_pre_nms_top_n_train: int = 1000,
    rpn_post_nms_top_n_train: int = 500,
    rpn_pre_nms_top_n_test: int = 500,
    rpn_post_nms_top_n_test: int = 200,
    box_batch_size_per_image: int = 128,
) -> torch.nn.Module:
    """Build Mask R-CNN with COCO-pretrained FPN/RPN/ROI heads + SSL backbone.

    Strategy:
    1. Create full COCO-pretrained Mask R-CNN detector
    2. Load SSL-pretrained ResNet50 state dict
    3. Replace only backbone.body weights with SSL weights
    4. Keep COCO-pretrained FPN / RPN / ROI box head / ROI mask head
    5. Replace box/mask predictors for LIVECell (2 classes)
    """
    model = maskrcnn_resnet50_fpn(
        weights=MaskRCNN_ResNet50_FPN_Weights.DEFAULT,
        weights_backbone=None,
        min_size=min_size,
        max_size=max_size,
        box_detections_per_img=box_detections_per_img,
        rpn_pre_nms_top_n_train=rpn_pre_nms_top_n_train,
        rpn_post_nms_top_n_train=rpn_post_nms_top_n_train,
        rpn_pre_nms_top_n_test=rpn_pre_nms_top_n_test,
        rpn_post_nms_top_n_test=rpn_post_nms_top_n_test,
        box_batch_size_per_image=box_batch_size_per_image,
    )
    _assert_matched_pretrained_backbone_policy(model)

    ssl_state = torch.load(str(ssl_backbone_path), map_location=device, weights_only=True)
    if "model" in ssl_state:
        ssl_state = ssl_state["model"]

    # Guard against full checkpoint
    if any(k.startswith("backbone.") for k in ssl_state.keys()):
        sample = list(ssl_state.keys())[:3]
        raise ValueError(
            f"Checkpoint '{ssl_backbone_path}' contains 'backbone.'-prefixed keys "
            f"(sample: {sample}). Use backbone_resnet50_ssl.pth instead."
        )

    # Map SSL standalone ResNet50 -> Mask R-CNN backbone.body.*
    mapped = {}
    skipped = []
    for k, v in ssl_state.items():
        if k.startswith("encoder.") or k.startswith("fc.") or k.startswith("projection."):
            skipped.append(k)
            continue
        if "num_batches_tracked" in k:
            continue
        mapped[f"backbone.body.{k}"] = v
    if skipped:
        print(f"Skipped {len(skipped)} non-backbone keys: {skipped[:3]}...")
    if not mapped:
        raise ValueError(f"No backbone weights were mapped from: {ssl_backbone_path}")

    # Load only backbone.body weights onto COCO-pretrained model
    missing, unexpected = model.load_state_dict(mapped, strict=False)
    print(f"Loaded {len(mapped)} SSL backbone keys onto COCO-pretrained heads")

    backbone_missing = [
        k for k in missing
        if k.startswith("backbone.body.") and not k.startswith("backbone.body.fc.")
    ]
    if backbone_missing:
        raise RuntimeError(
            f"{len(backbone_missing)} backbone.body.* keys MISSING! "
            f"Backbone may be COCO-init. Examples: {backbone_missing[:5]}"
        )
    print(f"Missing keys (expected, fc + predictor dim mismatch): {len(missing)}")

    if unexpected:
        raise RuntimeError(f"Unexpected keys: {unexpected[:5]}")

    # Replace predictors for LIVECell
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = model.roi_heads.mask_predictor.conv5_mask.out_channels
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)

    return model


def main():
    parser = argparse.ArgumentParser(description="Transfer SSL backbone to Mask R-CNN")
    parser.add_argument("--ssl-backbone", type=Path, required=True, help="SSL backbone .pth file")
    parser.add_argument("--output", type=Path, required=True, help="Output checkpoint for Mask R-CNN fine-tuning")
    parser.add_argument("--num-classes", type=int, default=2)
    args = parser.parse_args()

    device = "cpu"
    model = build_maskrcnn_with_ssl_backbone(args.ssl_backbone, args.num_classes, device=device)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"model": model.state_dict()}, args.output)
    print(f"Saved Mask R-CNN + SSL backbone: {args.output}")


if __name__ == "__main__":
    main()
