"""Step 5: Train YOLO12n 2-class checkbox detector on synthetic + real data.

Trains on the yolo_2class/ dataset built by build_yolo_dataset.py, exports
best weights to ONNX, and copies the final model to model/.

Supports resuming from the last checkpoint if training was interrupted.

Usage:
    python datasets/checkbox/tools/train.py
    python datasets/checkbox/tools/train.py --resume          # Resume interrupted run
    python datasets/checkbox/tools/train.py --epochs 300      # Override epochs
    python datasets/checkbox/tools/train.py --batch 4         # Smaller batch for low memory
"""

import argparse
import csv
import logging
import os
import shutil
from pathlib import Path
from typing import Any, Dict

logger = logging.getLogger(__name__)

ROOT = Path(__file__).resolve().parent.parent  # datasets/checkbox/
YOLO_2CLASS_DIR = ROOT / "yolo_2class"
MODEL_DIR = ROOT / "model"
RUNS_DIR = YOLO_2CLASS_DIR / "runs"

# Training defaults
DEFAULT_BASE_MODEL = "yolo12n.pt"
DEFAULT_IMGSZ = 1024
DEFAULT_EPOCHS = 200
DEFAULT_BATCH = 16
DEFAULT_PATIENCE = 50
DEFAULT_CLS_LOSS = 1.5
SAVE_PERIOD = 25
ONNX_OPSET = 18
ONNX_MODEL_NAME = "checkbox_yolo12n_2class.onnx"
YOLO_CLASS_NAMES = ["checkbox_checked", "checkbox_unchecked"]


def _detect_device() -> str:
    """Auto-detect best available training device: cuda -> mps -> cpu."""
    try:
        import torch
        if torch.cuda.is_available():
            logger.info("Detected CUDA device")
            return "0"
        if torch.backends.mps.is_available():
            logger.info("Detected MPS (Apple Silicon) device")
            return "mps"
    except Exception as e:
        logger.debug("Device auto-detection failed, falling back to cpu: %r", e)
    return "cpu"


def _find_resume_weights() -> str | None:
    """Find the last.pt checkpoint for resuming interrupted training."""
    # Fast path: expected location
    last_pt = RUNS_DIR / "train" / "weights" / "last.pt"
    if last_pt.exists():
        return str(last_pt)

    # Fallback: search for newest last.pt under runs/
    if RUNS_DIR.exists():
        candidates = sorted(
            RUNS_DIR.glob("**/weights/last.pt"),
            key=lambda p: p.stat().st_mtime,
            reverse=True,
        )
        if candidates:
            return str(candidates[0])

    return None


def _validate_onnx(onnx_path: str):
    """Sanity-check the exported ONNX model's input/output shapes."""
    try:
        import onnxruntime as ort
    except ImportError:
        logger.warning("onnxruntime not installed — skipping ONNX validation")
        return

    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    inp = sess.get_inputs()[0]
    out = sess.get_outputs()[0]
    logger.info("ONNX input:  %s %s", inp.name, inp.shape)
    logger.info("ONNX output: %s %s", out.name, out.shape)

    # Accept both (1, 4+C, N) and (1, N, 4+C) layouts
    if len(out.shape) != 3:
        logger.warning(
            "Unexpected ONNX output rank %d: %s; skipping strict validation",
            len(out.shape), out.shape,
        )
        return

    expected = 4 + len(YOLO_CLASS_NAMES)
    dim1 = out.shape[1]
    dim2 = out.shape[2]

    if isinstance(dim1, int) and dim1 == expected:
        num_classes = dim1 - 4
    elif isinstance(dim2, int) and dim2 == expected:
        num_classes = dim2 - 4
    else:
        logger.warning(
            "Unexpected ONNX output shape %s; skipping strict class-count check",
            out.shape,
        )
        return

    logger.info("ONNX validation passed (%d classes)", num_classes)


def _log_metrics(results_csv: Path) -> Dict[str, float]:
    """Read final metrics from results.csv."""
    metrics: Dict[str, float] = {}
    with open(results_csv, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        last = None
        for row in reader:
            last = row

    if not last:
        return metrics

    for key in last:
        clean = key.strip()
        if "mAP50" in clean and "mAP50-95" not in clean:
            try:
                metrics["mAP50"] = float(last[key])
            except (ValueError, TypeError):
                pass
        elif "mAP50-95" in clean:
            try:
                metrics["mAP50_95"] = float(last[key])
            except (ValueError, TypeError):
                pass

    return metrics


def _resolve_train_dir(model, results, resume_weights: str | None) -> Path:
    """Reliably determine the training output directory."""
    # Try results.save_dir first
    save_dir = getattr(results, "save_dir", None)
    if save_dir:
        return Path(save_dir)

    # Try model.trainer.save_dir
    trainer = getattr(model, "trainer", None)
    save_dir = getattr(trainer, "save_dir", None)
    if save_dir:
        return Path(save_dir)

    # Derive from resume checkpoint path: .../train/weights/last.pt -> .../train
    if resume_weights:
        return Path(resume_weights).resolve().parents[1]

    raise RuntimeError(
        "Could not determine training save_dir. "
        "Check Ultralytics version compatibility."
    )


def train(
    resume: bool = False,
    epochs: int = DEFAULT_EPOCHS,
    batch: int = DEFAULT_BATCH,
    imgsz: int = DEFAULT_IMGSZ,
    base_model: str = DEFAULT_BASE_MODEL,
) -> Dict[str, Any]:
    """Train the 2-class checkbox detector.

    Args:
        resume: If True, resume from last checkpoint.
        epochs: Max training epochs.
        batch: Batch size.
        imgsz: Input image size.
        base_model: Pretrained model to start from.

    Returns:
        Dict with training stats and output paths.
    """
    try:
        from ultralytics import YOLO
    except ImportError:
        raise ImportError(
            "ultralytics is required for training. "
            "Install with: pip install ultralytics"
        )

    dataset_yaml = YOLO_2CLASS_DIR / "dataset.yaml"
    if not dataset_yaml.exists():
        raise FileNotFoundError(
            f"Dataset YAML not found: {dataset_yaml}\n"
            "Run build_yolo_dataset.py first."
        )

    stats: Dict[str, Any] = {}
    device = _detect_device()
    resume_weights: str | None = None

    # ── Resume or fresh start ─────────────────────────────────────────────
    if resume:
        resume_weights = _find_resume_weights()
        if resume_weights is None:
            logger.warning("No checkpoint found to resume from. Starting fresh.")
            resume = False
        else:
            logger.info("Resuming from checkpoint: %s", resume_weights)
    else:
        # Guard against accidental overwrite of existing run
        existing_last = RUNS_DIR / "train" / "weights" / "last.pt"
        if existing_last.exists():
            raise RuntimeError(
                f"Found existing checkpoint at {existing_last}. "
                "Use --resume to continue, or delete runs/train/ to start fresh."
            )

    if resume:
        assert resume_weights is not None  # guaranteed by logic above
        model = YOLO(resume_weights)
        logger.info("Resuming training on device=%s ...", device)
        results = model.train(
            resume=True,
            device=device,
        )
    else:
        model = YOLO(base_model)
        logger.info(
            "Starting training: model=%s, device=%s, imgsz=%d, batch=%d, epochs=%d",
            base_model, device, imgsz, batch, epochs,
        )
        # Mosaic is device-conditional: safe on CUDA, can OOM on MPS
        # with dense pages (100+ checkboxes). close_mosaic disables it
        # for the last N epochs to stabilize fine-tuning.
        is_cuda = device not in ("cpu", "mps")
        mosaic = 1.0 if is_cuda else 0.0

        results = model.train(
            data=str(dataset_yaml),
            epochs=epochs,
            imgsz=imgsz,
            batch=batch,
            device=device,
            patience=DEFAULT_PATIENCE,
            cos_lr=True,
            cls=DEFAULT_CLS_LOSS,
            save_period=SAVE_PERIOD,
            project=str(RUNS_DIR),
            name="train",
            exist_ok=True,
            verbose=True,
            mosaic=mosaic,
            close_mosaic=10 if mosaic else 0,
            cache="ram" if is_cuda else False,
            workers=min(8, os.cpu_count() or 2),
        )

    # ── Find best weights ─────────────────────────────────────────────────
    train_dir = _resolve_train_dir(model, results, resume_weights)
    best_pt = train_dir / "weights" / "best.pt"
    if not best_pt.exists():
        raise FileNotFoundError(
            f"Training did not produce best.pt at {best_pt}. "
            "Check training logs above."
        )

    stats["best_pt"] = str(best_pt)
    logger.info("Training complete. Best weights: %s", best_pt)

    # ── Read metrics ──────────────────────────────────────────────────────
    results_csv = train_dir / "results.csv"
    if results_csv.exists():
        metrics = _log_metrics(results_csv)
        stats.update(metrics)
        if "mAP50" in metrics:
            logger.info("Final mAP50: %.4f", metrics["mAP50"])
        if "mAP50_95" in metrics:
            logger.info("Final mAP50-95: %.4f", metrics["mAP50_95"])

    # ── Export to ONNX ────────────────────────────────────────────────────
    # Use the training run's imgsz when resuming to avoid mismatch
    trainer = getattr(model, "trainer", None)
    resume_imgsz = getattr(getattr(trainer, "args", None), "imgsz", None)
    export_imgsz = resume_imgsz if (resume and resume_imgsz) else imgsz

    logger.info("Exporting to ONNX (imgsz=%s)...", export_imgsz)
    best_model = YOLO(str(best_pt))
    export_path = best_model.export(
        format="onnx",
        imgsz=export_imgsz,
        simplify=True,
        opset=ONNX_OPSET,
    )

    # Validate export result
    if isinstance(export_path, (list, tuple)):
        export_path = export_path[0]
    if not export_path or not Path(export_path).exists():
        raise FileNotFoundError(
            f"ONNX export failed or returned invalid path: {export_path}"
        )

    logger.info("ONNX exported: %s", export_path)
    _validate_onnx(str(export_path))

    # ── Copy to model/ ────────────────────────────────────────────────────
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    final_path = MODEL_DIR / ONNX_MODEL_NAME
    shutil.copy2(str(export_path), str(final_path))
    stats["onnx_path"] = str(final_path)
    logger.info("Final model: %s", final_path)

    return stats


def main():
    parser = argparse.ArgumentParser(
        description="Train YOLO12n 2-class checkbox detector"
    )
    parser.add_argument(
        "--resume", action="store_true",
        help="Resume training from last checkpoint",
    )
    parser.add_argument("--epochs", type=int, default=DEFAULT_EPOCHS)
    parser.add_argument("--batch", type=int, default=DEFAULT_BATCH)
    parser.add_argument("--imgsz", type=int, default=DEFAULT_IMGSZ)
    parser.add_argument("--base-model", type=str, default=DEFAULT_BASE_MODEL)
    args = parser.parse_args()

    stats = train(
        resume=args.resume,
        epochs=args.epochs,
        batch=args.batch,
        imgsz=args.imgsz,
        base_model=args.base_model,
    )

    logger.info("\n=== Training Summary ===")
    for k, v in sorted(stats.items()):
        logger.info("  %s: %s", k, v)


if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s: %(message)s",
        datefmt="%H:%M:%S",
    )
    main()
