#!/usr/bin/env python3
"""Fine-tune Mask R-CNN with SSL-pretrained backbone on LIVECell.

Reuses the training and evaluation logic from train_maskrcnn_baseline.py,
but initializes the backbone from an SSL checkpoint instead of COCO weights.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch
from torch.utils.data import DataLoader

# Reuse baseline utilities
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from train_maskrcnn_baseline import (
    LiveCellCocoDataset,
    discover_size_split,
    discover_val_ann_file,
    evaluate_and_visualize,
    index_images,
    train_model,
    write_csv,
)
from adaptation_optimizer import (
    build_adaptation_optimizer,
    optimizer_group_summary,
    sha256_file,
    snapshot_trainable_parameters,
    validate_multiplier,
    write_parameter_adaptation,
)
from livecell_protocol import apply_livecell_dense_protocol, dense_protocol_manifest
from train_maskrcnn_smoke import collate_fn, set_seed

# Import SSL backbone transfer
sys.path.insert(0, str(Path(__file__).resolve().parent))
from transfer_backbone import build_maskrcnn_with_coco_heads_ssl_backbone, build_maskrcnn_with_ssl_backbone


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--data-root", type=Path, required=True)
    p.add_argument("--ssl-backbone", type=Path, required=True, help="SSL backbone .pth")
    p.add_argument("--output-dir", type=Path, required=True)
    p.add_argument("--train-ann-file", type=Path, default=None,
                   help="Override train COCO JSON (for cross cell-line splits)")
    p.add_argument("--val-ann-file", type=Path, default=None,
                   help="Override val COCO JSON (for cross cell-line splits)")
    p.add_argument("--train-size-split", type=str, default="5", choices=["2", "4", "5", "25", "50"])
    p.add_argument("--epochs", type=int, default=20)
    p.add_argument("--max-train-images", type=int, default=0)
    p.add_argument("--max-val-images", type=int, default=0)
    p.add_argument("--max-instances-per-image", type=int, default=128)
    p.add_argument("--livecell-dense-protocol", action="store_true")
    p.add_argument("--train-crop-size", type=int, default=-1)
    p.add_argument("--batch-size", type=int, default=1)
    p.add_argument("--num-workers", type=int, default=2)
    p.add_argument("--lr", type=float, default=0.0001)
    p.add_argument("--min-size", type=int, default=512)
    p.add_argument("--max-size", type=int, default=768)
    p.add_argument("--eval-score-threshold", type=float, default=0.05)
    p.add_argument("--eval-max-predictions", type=int, default=100)
    p.add_argument("--box-detections-per-image", type=int, default=100)
    p.add_argument("--rpn-pre-nms-top-n-train", type=int, default=1000)
    p.add_argument("--rpn-post-nms-top-n-train", type=int, default=500)
    p.add_argument("--rpn-pre-nms-top-n-test", type=int, default=500)
    p.add_argument("--rpn-post-nms-top-n-test", type=int, default=200)
    p.add_argument("--box-batch-size-per-image", type=int, default=128)
    p.add_argument("--eval-max-images", type=int, default=0)
    p.add_argument("--vis-count", type=int, default=8)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--device", default="auto")
    p.add_argument("--defer-coco-eval", action="store_true",
                   help="Skip built-in COCOeval; write predictions only for external protocol ablation.")
    p.add_argument("--hybrid-coco-heads", action="store_true",
                   help="Use COCO-pretrained FPN/RPN/ROI heads + SSL backbone (hybrid ablation)")
    p.add_argument("--max-steps", type=int, default=0,
                   help="0 means no explicit step limit.")
    p.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 p.parse_args()


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 = args.train_ann_file or discover_size_split(args.data_root, args.train_size_split)
    val_ann = args.val_ann_file or discover_val_ann_file(args.data_root)

    print(f"DEVICE: {device}")
    print(f"TRAIN: {train_ann}")
    print(f"VAL: {val_ann}")
    print(f"SSL backbone: {args.ssl_backbone}")

    image_index = index_images(args.data_root)
    train_dataset = LiveCellCocoDataset(
        args.data_root, train_ann, 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, 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",
    )

    # Build model
    if args.hybrid_coco_heads:
        print("HYBRID MODE: COCO-pretrained FPN/RPN/ROI heads + SSL backbone")
        model = build_maskrcnn_with_coco_heads_ssl_backbone(
            args.ssl_backbone,
            num_classes=2,
            min_size=args.min_size,
            max_size=args.max_size,
            device=str(device),
            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,
        )
    else:
        model = build_maskrcnn_with_ssl_backbone(
            args.ssl_backbone,
            num_classes=2,
            min_size=args.min_size,
            max_size=args.max_size,
            device=str(device),
            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)

    optimizer = build_adaptation_optimizer(
        model,
        base_lr=args.lr,
        multiplier=args.backbone_lr_multiplier,
        momentum=0.9,
        weight_decay=0.0005,
    )
    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()},
        "ssl_backbone": str(args.ssl_backbone),
        "train_size_split": args.train_size_split,
        "hybrid_coco_heads": args.hybrid_coco_heads,
        "device": str(device),
        "torch": torch.__version__,
        "protocol": dense_protocol_manifest(args),
        "train_dataset": train_dataset.describe(),
        "val_dataset": val_dataset.describe(),
        "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")

    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)

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

    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(ckpt),
        "checkpoint_sha256": sha256_file(ckpt),
        "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("SSL_FINETUNE_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(ckpt),
        }
        (args.output_dir / "inference_complete.json").write_text(json.dumps(inference_complete, indent=2))


if __name__ == "__main__":
    main()
