src/segmentation/trainer.py

"""
trainer.py
----------
YOLOv8 segmentation trainer for floor plan element detection.

Automatically selects the best available device:
  - Apple Silicon (M1/M2/M3/M4): uses MPS
  - NVIDIA GPU: uses CUDA
  - Fallback: CPU

Usage:
    from src.segmentation.trainer import SegmentationTrainer

    trainer = SegmentationTrainer(dataset_yaml="data/yolo_dataset/dataset.yaml")
    trainer.train()
    trainer.export()
"""

import os
import platform
from pathlib import Path
from typing import Optional

import torch


# ── Device detection ──────────────────────────────────────────────────────────

def get_best_device() -> str:
    """
    Returns the best available device string for PyTorch / Ultralytics.
    Priority: MPS (Apple Silicon) > CUDA (if compatible) > CPU
    """
    if torch.backends.mps.is_available():
        print("Device: Apple Silicon MPS (GPU accelerated)")
        return "mps"
    elif torch.cuda.is_available():
        # Check CUDA compute capability — RTX 5060 (sm_120) not supported by older PyTorch
        major, minor = torch.cuda.get_device_capability(0)
        sm = major * 10 + minor
        supported = [50, 60, 61, 70, 75, 80, 86, 90]
        if sm in supported:
            gpu = torch.cuda.get_device_name(0)
            print(f"Device: CUDA — {gpu}")
            return "0"
        else:
            gpu = torch.cuda.get_device_name(0)
            print(f"Device: CPU (GPU {gpu} sm_{sm} not supported by this PyTorch build)")
            return "cpu"
    else:
        print("Device: CPU (training will be slow — consider Colab)")
        return "cpu"


# ── Trainer ───────────────────────────────────────────────────────────────────

class SegmentationTrainer:
    """
    Wraps Ultralytics YOLOv8 for floor plan segmentation training.

    Args:
        dataset_yaml:  Path to the dataset.yaml generated by CubiCasaDataset.
        model_size:    YOLOv8 model variant: 'n' (nano), 's' (small), 'm' (medium).
                       For M4 MacBook Air, 's' is the sweet spot.
        output_dir:    Where to save training runs.
        epochs:        Number of training epochs.
        batch_size:    Batch size. Use 8-16 for MPS, 4-8 for CPU.
        img_size:      Input image size for training.
        device:        Override device ('mps', 'cpu', '0'). Auto-detected if None.
    """

    # Model weight files (downloaded automatically on first run)
    MODEL_WEIGHTS = {
        "n": "yolov8n-seg.pt",   # 3.4M params — fastest
        "s": "yolov8s-seg.pt",   # 11.8M params — recommended for M4
        "m": "yolov8m-seg.pt",   # 27.3M params — most accurate
    }

    def __init__(
        self,
        dataset_yaml: str,
        model_size: str = "s",
        output_dir: str = "models/segmentation",
        epochs: int = 100,
        batch_size: int = 8,
        img_size: int = 640,
        device: Optional[str] = None,
        pretrained: bool = True,
    ):
        self.dataset_yaml = Path(dataset_yaml)
        self.model_size = model_size
        self.output_dir = Path(output_dir)
        self.epochs = epochs
        self.batch_size = batch_size
        self.img_size = img_size
        self.device = device or get_best_device()
        self.pretrained = pretrained
        self.model = None

        if model_size not in self.MODEL_WEIGHTS:
            raise ValueError(
                f"model_size must be one of {list(self.MODEL_WEIGHTS.keys())}"
            )
        if not self.dataset_yaml.exists():
            raise FileNotFoundError(f"Dataset YAML not found: {dataset_yaml}")

    def train(self) -> str:
        """
        Start YOLOv8 training.

        Returns:
            Path to the best model weights file.
        """
        from ultralytics import YOLO

        weights = self.MODEL_WEIGHTS[self.model_size]
        print(f"\nStarting training:")
        print(f"  Model:   YOLOv8{self.model_size}-seg ({weights})")
        print(f"  Dataset: {self.dataset_yaml}")
        print(f"  Device:  {self.device}")
        print(f"  Epochs:  {self.epochs}")
        print(f"  Batch:   {self.batch_size}")
        print(f"  Img size:{self.img_size}")
        print()

        self.model = YOLO(weights)

        results = self.model.train(
            data=str(self.dataset_yaml),
            epochs=self.epochs,
            batch=self.batch_size,
            imgsz=self.img_size,
            device=self.device,
            project=str(self.output_dir),
            name="floor_plan_seg",
            # Augmentation — essential for floor plans
            augment=True,
            degrees=10.0,       # Random rotation ±10°
            translate=0.1,      # Random translation
            scale=0.5,          # Random scale
            fliplr=0.5,         # Horizontal flip
            flipud=0.0,         # No vertical flip (plans are top-down)
            # Optimizer settings
            optimizer="AdamW",
            lr0=0.001,
            lrf=0.01,
            weight_decay=0.0005,
            warmup_epochs=3,
            # Save settings
            save=True,
            save_period=10,     # Save checkpoint every 10 epochs
            patience=30,        # Early stopping patience
            # Logging
            verbose=True,
            plots=True,
        )

        best_weights = (
            Path(results.save_dir) / "weights" / "best.pt"
        )
        print(f"\nTraining complete. Best weights: {best_weights}")
        return str(best_weights)

    def resume(self, checkpoint_path: str) -> str:
        """Resume training from a checkpoint."""
        from ultralytics import YOLO
        print(f"Resuming from checkpoint: {checkpoint_path}")
        self.model = YOLO(checkpoint_path)
        results = self.model.train(resume=True)
        return str(Path(results.save_dir) / "weights" / "best.pt")

    def validate(self, weights_path: Optional[str] = None) -> dict:
        """
        Run validation on the test set.

        Returns:
            Dictionary of evaluation metrics (mAP50, mAP50-95, etc.)
        """
        from ultralytics import YOLO

        if weights_path:
            model = YOLO(weights_path)
        elif self.model:
            model = self.model
        else:
            raise ValueError(
                "No model loaded. Run train() first or pass weights_path."
            )

        metrics = model.val(
            data=str(self.dataset_yaml),
            device=self.device,
            split="test",
        )

        print("\nValidation results:")
        print(f"  mAP50:    {metrics.seg.map50:.4f}")
        print(f"  mAP50-95: {metrics.seg.map:.4f}")
        return {
            "map50":    metrics.seg.map50,
            "map50_95": metrics.seg.map,
        }

    def export(
        self,
        weights_path: Optional[str] = None,
        format: str = "onnx",
    ) -> str:
        """
        Export trained model for production deployment.

        Args:
            weights_path: Path to best.pt. Uses last trained model if None.
            format:       Export format: 'onnx', 'coreml' (Apple), 'torchscript'.

        Returns:
            Path to exported model file.
        """
        from ultralytics import YOLO

        if weights_path:
            model = YOLO(weights_path)
        elif self.model:
            model = self.model
        else:
            raise ValueError("No model to export. Run train() first.")

        print(f"Exporting model to {format.upper()} format...")

        # CoreML is best for on-device M4 inference
        if format == "coreml":
            path = model.export(
                format="coreml",
                imgsz=self.img_size,
                nms=True,
            )
        else:
            path = model.export(
                format=format,
                imgsz=self.img_size,
                dynamic=True,
                simplify=True,
            )

        print(f"Exported: {path}")
        return str(path)