#!/usr/bin/env python3
"""Run a tiny LIVECell Mask R-CNN training smoke test.

This is not the final experiment script. It is a deliberately small baseline
used to verify that LIVECell loading, instance masks, GPU training, checkpoints,
and prediction visualization all work on the remote machine.
"""

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, Iterable, List, Optional, Sequence, Tuple

import numpy as np
import torch
from PIL import Image, ImageDraw
from pycocotools import mask as mask_utils
from torch.utils.data import DataLoader, Dataset
from torchvision.models import ResNet50_Weights, resnet50
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


IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-root", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, default=Path("results/maskrcnn_smoke"))
    parser.add_argument("--train-split", default="train")
    parser.add_argument("--val-split", default="val")
    parser.add_argument("--max-train-images", type=int, default=8)
    parser.add_argument("--max-val-images", type=int, default=4)
    parser.add_argument("--max-instances-per-image", type=int, default=64)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--max-steps", type=int, default=8)
    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=384)
    parser.add_argument("--max-size", type=int, default=640)
    parser.add_argument("--vis-count", type=int, default=4)
    parser.add_argument("--score-threshold", type=float, default=0.05)
    parser.add_argument("--max-predictions", type=int, default=30)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--pretrained-coco",
        action="store_true",
        help="Use torchvision COCO weights. Omit this on fresh machines to avoid downloads.",
    )
    parser.add_argument(
        "--device",
        default="auto",
        choices=["auto", "cuda", "cpu"],
        help="Device for training. 'auto' uses CUDA when available.",
    )
    return parser.parse_args()


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def discover_annotation_file(data_root: Path, split: str) -> Path:
    json_files = sorted(data_root.rglob("*.json"))
    candidates = [
        p
        for p in json_files
        if split.lower() in p.name.lower()
        and ("coco" in p.name.lower() or "annotation" in p.name.lower())
    ]
    if not candidates:
        candidates = [p for p in json_files if split.lower() in p.name.lower()]
    if not candidates:
        raise FileNotFoundError(f"Could not find a COCO annotation file for split={split!r} under {data_root}")
    return candidates[0]


def index_images(data_root: Path) -> Dict[str, List[Path]]:
    image_index: Dict[str, List[Path]] = defaultdict(list)
    for path in data_root.rglob("*"):
        if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES:
            image_index[path.name].append(path)
    return image_index


def find_image_path(file_name: str, image_index: Dict[str, List[Path]], data_root: Path) -> Optional[Path]:
    direct = data_root / file_name
    if direct.exists():
        return direct

    by_name = image_index.get(Path(file_name).name, [])
    if by_name:
        return by_name[0]
    return None


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


def annotation_to_mask(annotation: Dict[str, Any], height: int, width: int) -> np.ndarray:
    segmentation = annotation.get("segmentation")
    if segmentation is None:
        return np.zeros((height, width), dtype=np.uint8)

    if isinstance(segmentation, list):
        rles = mask_utils.frPyObjects(segmentation, height, width)
        rle = mask_utils.merge(rles)
    elif isinstance(segmentation, dict):
        if isinstance(segmentation.get("counts"), list):
            rle = mask_utils.frPyObjects(segmentation, height, width)
        else:
            rle = segmentation
    else:
        return np.zeros((height, width), dtype=np.uint8)

    decoded = mask_utils.decode(rle)
    if decoded.ndim == 3:
        decoded = np.max(decoded, axis=2)
    return decoded.astype(np.uint8)


def pil_to_tensor(image: Image.Image) -> torch.Tensor:
    array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
    return torch.from_numpy(array).permute(2, 0, 1).contiguous()


class LiveCellInstanceDataset(Dataset):
    def __init__(
        self,
        data_root: Path,
        split: str,
        image_index: Dict[str, List[Path]],
        max_images: int,
        max_instances_per_image: int,
        seed: int,
    ) -> None:
        self.data_root = data_root
        self.split = split
        self.annotation_file = discover_annotation_file(data_root, split)
        self.coco = load_coco(self.annotation_file)
        self.max_instances_per_image = max_instances_per_image

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

        rng = random.Random(seed)
        items = []
        for img in self.coco.get("images", []):
            image_id = int(img["id"])
            anns = anns_by_image.get(image_id, [])
            if not anns:
                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),
                }
            )

        # Small-instance images reduce memory pressure during the smoke test.
        items.sort(key=lambda item: item["annotation_count"])
        if max_images > 0:
            pool = items[: max(max_images * 8, max_images)]
            rng.shuffle(pool)
            items = sorted(pool[:max_images], key=lambda item: item["annotation_count"])
        self.items = items

        if not self.items:
            raise RuntimeError(f"No usable images found for split={split!r}")

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

    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"]

        with Image.open(image_path) as image:
            image = image.convert("RGB")
            width, height = image.size
            image_tensor = pil_to_tensor(image)

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

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

        for ann in annotations:
            bbox = ann.get("bbox")
            if not bbox or len(bbox) != 4:
                continue
            x, y, w, h = [float(v) for v in bbox]
            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

            mask = annotation_to_mask(ann, height, width)
            if not mask.any():
                continue

            boxes.append([x1, y1, x2, y2])
            masks.append(mask)
            areas.append(float(ann.get("area", w * h)))
            iscrowd.append(int(ann.get("iscrowd", 0)))

        if not boxes:
            # Very rare for valid LIVECell entries, but keeps the model API sane.
            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 {
            "split": self.split,
            "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)),
            "max_instances_per_image_used": self.max_instances_per_image,
        }


def collate_fn(batch: Sequence[Tuple[torch.Tensor, Dict[str, torch.Tensor]]]) -> Tuple[List[torch.Tensor], List[Dict[str, torch.Tensor]]]:
    images, targets = zip(*batch)
    return list(images), list(targets)


def build_model(num_classes: int, min_size: int, max_size: int,
                pretrained_coco: bool = False, imagenet_backbone: bool = False,
                imagenet_backbone_coco_heads: bool = False,
                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:
    flags = [pretrained_coco, imagenet_backbone, imagenet_backbone_coco_heads]
    if sum(bool(f) for f in flags) > 1:
        raise ValueError(
            "pretrained_coco / imagenet_backbone / imagenet_backbone_coco_heads "
            "are mutually exclusive"
        )
    if pretrained_coco:
        weights = "DEFAULT"
        weights_backbone = None
    elif imagenet_backbone:
        weights = None
        weights_backbone = ResNet50_Weights.IMAGENET1K_V2
    elif imagenet_backbone_coco_heads:
        weights = MaskRCNN_ResNet50_FPN_Weights.DEFAULT
        weights_backbone = None
    else:
        weights = None
        weights_backbone = None
    kwargs = {
        "weights": weights,
        "weights_backbone": weights_backbone,
        "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,
    }
    model = maskrcnn_resnet50_fpn(**kwargs)

    # ImageNet V2 backbone + COCO heads: replace backbone.body manually
    if imagenet_backbone_coco_heads:
        imgnet_resnet = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
        state = imgnet_resnet.state_dict()
        mapped = {}
        for k, v in state.items():
            if k.startswith("fc.") or "num_batches_tracked" in k:
                continue
            mapped[f"backbone.body.{k}"] = v
        missing, unexpected = model.load_state_dict(mapped, strict=False)
        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"Examples: {backbone_missing[:5]}"
            )
        if unexpected:
            raise RuntimeError(f"Unexpected keys: {unexpected[:5]}")
        print(f"Loaded {len(mapped)} ImageNet V2 backbone keys onto COCO heads")

    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 move_targets_to_device(targets: List[Dict[str, torch.Tensor]], device: torch.device) -> List[Dict[str, torch.Tensor]]:
    return [{k: v.to(device) for k, v in target.items()} for target in targets]


def train_one_smoke_run(
    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,
                "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 draw_masks(
    image_array: np.ndarray,
    masks: np.ndarray,
    boxes: np.ndarray,
    scores: Optional[np.ndarray],
    output_path: Path,
    max_items: int,
) -> int:
    overlay = image_array.astype(np.float32).copy()
    rng = random.Random(123)
    count = min(len(masks), max_items)

    for idx in range(count):
        mask = masks[idx].astype(bool)
        if not mask.any():
            continue
        color = np.array([rng.randint(40, 255), rng.randint(40, 255), rng.randint(40, 255)], dtype=np.float32)
        overlay[mask] = overlay[mask] * 0.55 + color * 0.45

    result = Image.fromarray(np.clip(overlay, 0, 255).astype(np.uint8))
    draw = ImageDraw.Draw(result)
    for idx in range(count):
        x1, y1, x2, y2 = boxes[idx].tolist()
        draw.rectangle([x1, y1, x2, y2], outline=(255, 255, 0), width=1)
        if scores is not None:
            draw.text((x1, max(0, y1 - 10)), f"{scores[idx]:.2f}", fill=(255, 255, 0))

    output_path.parent.mkdir(parents=True, exist_ok=True)
    result.save(output_path)
    return count


@torch.no_grad()
def visualize_predictions(
    model: torch.nn.Module,
    dataset: LiveCellInstanceDataset,
    device: torch.device,
    output_dir: Path,
    vis_count: int,
    score_threshold: float,
    max_predictions: int,
) -> List[Dict[str, Any]]:
    model.eval()
    rows: List[Dict[str, Any]] = []

    for index in range(min(vis_count, len(dataset))):
        image, target = dataset[index]
        prediction = model([image.to(device)])[0]
        image_array = (image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)

        gt_boxes = target["boxes"].cpu().numpy()
        gt_masks = target["masks"].cpu().numpy()
        gt_count = draw_masks(
            image_array,
            gt_masks,
            gt_boxes,
            None,
            output_dir / "visualization" / f"val_{index:03d}_gt.png",
            max_items=max_predictions,
        )

        scores = prediction["scores"].detach().cpu()
        keep = scores >= score_threshold
        pred_scores = scores[keep].numpy()
        pred_boxes = prediction["boxes"].detach().cpu()[keep].numpy()
        pred_masks = (prediction["masks"].detach().cpu()[keep, 0].numpy() >= 0.5).astype(np.uint8)
        pred_count = draw_masks(
            image_array,
            pred_masks,
            pred_boxes,
            pred_scores,
            output_dir / "visualization" / f"val_{index:03d}_pred.png",
            max_items=max_predictions,
        )

        rows.append(
            {
                "val_index": index,
                "image_id": int(target["image_id"].item()),
                "gt_instances_drawn": gt_count,
                "predictions_above_threshold": int(keep.sum().item()),
                "predictions_drawn": pred_count,
                "top_score": float(scores[0].item()) if len(scores) else 0.0,
            }
        )

    return rows


def write_csv(path: Path, rows: List[Dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if not rows:
        path.write_text("", encoding="utf-8")
        return
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    args = 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)

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

    image_index = index_images(args.data_root)
    train_dataset = LiveCellInstanceDataset(
        args.data_root,
        args.train_split,
        image_index,
        args.max_train_images,
        args.max_instances_per_image,
        args.seed,
    )
    val_dataset = LiveCellInstanceDataset(
        args.data_root,
        args.val_split,
        image_index,
        args.max_val_images,
        args.max_instances_per_image,
        args.seed + 1,
    )

    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",
    )

    model = build_model(num_classes=2, min_size=args.min_size, max_size=args.max_size, pretrained_coco=args.pretrained_coco)
    model.to(device)

    optimizer = torch.optim.SGD(
        [param for param in model.parameters() if param.requires_grad],
        lr=args.lr,
        momentum=0.9,
        weight_decay=0.0005,
    )

    config = {
        "args": {key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()},
        "device": str(device),
        "torch": torch.__version__,
        "torchvision_model": "maskrcnn_resnet50_fpn",
        "train_dataset": train_dataset.describe(),
        "val_dataset": val_dataset.describe(),
    }
    (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_one_smoke_run(model, train_loader, optimizer, device, args.epochs, args.max_steps)
    write_csv(args.output_dir / "loss_log.csv", loss_rows)

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

    pred_rows = visualize_predictions(
        model,
        val_dataset,
        device,
        args.output_dir,
        args.vis_count,
        args.score_threshold,
        args.max_predictions,
    )
    write_csv(args.output_dir / "quick_eval.csv", pred_rows)

    final_summary = {
        "status": "ok",
        "steps": len(loss_rows),
        "last_total_loss": loss_rows[-1]["total_loss"] if loss_rows else None,
        "checkpoint": str(checkpoint_path),
        "loss_log": str(args.output_dir / "loss_log.csv"),
        "quick_eval": str(args.output_dir / "quick_eval.csv"),
        "visualization_dir": str(args.output_dir / "visualization"),
    }
    (args.output_dir / "summary.json").write_text(json.dumps(final_summary, indent=2), encoding="utf-8")
    print("SMOKE_SUMMARY", json.dumps(final_summary, ensure_ascii=False))


if __name__ == "__main__":
    main()
