"""Step 8: Train YOLO12 checkbox detector and export to ONNX.

Converts COCO annotations → YOLO format, trains with Ultralytics,
and exports best.pt → ONNX for use with onnx_engine.py.

Usage:
    python -m datasets.checkbox.pipeline.run train
"""

import json
import logging
import shutil
from pathlib import Path
from typing import Any, Dict

from .config import (
    ANNOTATIONS_DIR,
    CHECKED_OVERSAMPLE,
    COCO_TO_YOLO_CLASS,
    IMAGES_DIR,
    MODEL_DIR,
    ONNX_MODEL_NAME,
    ONNX_OPSET,
    YOLO_BASE_MODEL,
    YOLO_BATCH,
    YOLO_CLASS_NAMES,
    YOLO_CLS_LOSS,
    YOLO_DIR,
    YOLO_EPOCHS,
    YOLO_IMGSZ,
    YOLO_PATIENCE,
    YOLO_SAVE_PERIOD,
)

logger = logging.getLogger(__name__)


def _detect_device() -> str:
    """Auto-detect best available training device: mps → cuda → cpu."""
    try:
        import torch

        if torch.backends.mps.is_available():
            return "mps"
        if torch.cuda.is_available():
            return "0"  # first CUDA GPU
    except Exception:
        pass
    return "cpu"


def _clean_yolo_dirs():
    """Remove stale YOLO output dirs so re-runs start fresh."""
    for subdir in ("train", "val"):
        d = YOLO_DIR / subdir
        if d.exists():
            shutil.rmtree(d)
            logger.info("Cleaned stale YOLO dir: %s", d)


def _convert_coco_to_yolo(split: str) -> int:
    """Convert one COCO split to YOLO label files.

    Reads annotations/{split}.json and creates:
      yolo/{split}/images/  (symlinks to originals)
      yolo/{split}/labels/  (one .txt per image)

    For the train split, images containing checked checkboxes are
    duplicated CHECKED_OVERSAMPLE times to mitigate class imbalance.

    Returns number of images processed (including oversampled copies).
    """
    coco_path = ANNOTATIONS_DIR / f"{split}.json"
    if not coco_path.exists():
        raise FileNotFoundError(f"COCO annotations not found: {coco_path}")

    with open(coco_path) as f:
        coco = json.load(f)

    # Build lookup: image_id → image info
    images = {img["id"]: img for img in coco["images"]}

    # Group annotations by image_id
    annots_by_image: Dict[int, list] = {}
    for annot in coco["annotations"]:
        annots_by_image.setdefault(annot["image_id"], []).append(annot)

    # Identify which COCO category IDs map to "checked" (YOLO class 0)
    checked_coco_ids = {
        coco_id for coco_id, yolo_cls in COCO_TO_YOLO_CLASS.items() if yolo_cls == 0
    }

    # Output dirs
    img_out = YOLO_DIR / split / "images"
    lbl_out = YOLO_DIR / split / "labels"
    img_out.mkdir(parents=True, exist_ok=True)
    lbl_out.mkdir(parents=True, exist_ok=True)

    count = 0
    for img_id, img_info in images.items():
        filename = img_info["file_name"]
        img_w = img_info["width"]
        img_h = img_info["height"]

        src = IMAGES_DIR / split / filename
        if not src.exists():
            logger.warning("Image not found: %s", src)
            continue

        # Build label lines for this image
        img_annots = annots_by_image.get(img_id, [])
        lines = []
        has_checked = False

        for annot in img_annots:
            coco_cat_id = annot["category_id"]
            yolo_cls = COCO_TO_YOLO_CLASS.get(coco_cat_id)
            if yolo_cls is None:
                logger.warning("Unknown COCO category %d, skipping", coco_cat_id)
                continue

            if coco_cat_id in checked_coco_ids:
                has_checked = True

            # COCO bbox: [x, y, w, h] in pixels → clamp in pixel space first
            bx, by, bw, bh = annot["bbox"]
            x0 = max(0.0, bx)
            y0 = max(0.0, by)
            x1 = min(float(img_w), bx + bw)
            y1 = min(float(img_h), by + bh)

            # Convert to YOLO normalized cx cy w h
            cx = ((x0 + x1) / 2) / img_w
            cy = ((y0 + y1) / 2) / img_h
            nw = (x1 - x0) / img_w
            nh = (y1 - y0) / img_h

            lines.append(f"{yolo_cls} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}")

        label_text = "\n".join(lines) + "\n" if lines else ""

        # Determine how many copies to write (oversampling for train split)
        copies = 1
        if split == "train" and has_checked and CHECKED_OVERSAMPLE > 1:
            copies = CHECKED_OVERSAMPLE

        for copy_idx in range(copies):
            if copy_idx == 0:
                # Original name
                dst_name = filename
                lbl_name = Path(filename).stem + ".txt"
            else:
                # Oversampled copies get a suffix
                stem = Path(filename).stem
                ext = Path(filename).suffix
                dst_name = f"{stem}_os{copy_idx}{ext}"
                lbl_name = f"{stem}_os{copy_idx}.txt"

            # Symlink image (remove stale symlink if present)
            dst = img_out / dst_name
            if dst.exists() or dst.is_symlink():
                dst.unlink()
            dst.symlink_to(src.resolve())

            # Write label file
            (lbl_out / lbl_name).write_text(label_text)
            count += 1

    return count


def _write_dataset_yaml() -> Path:
    """Write dataset.yaml for Ultralytics training. Returns path."""
    yaml_path = YOLO_DIR / "dataset.yaml"

    # Use absolute paths for reliability
    train_path = (YOLO_DIR / "train" / "images").resolve()
    val_path = (YOLO_DIR / "val" / "images").resolve()

    content = (
        f"path: {YOLO_DIR.resolve()}\n"
        f"train: {train_path}\n"
        f"val: {val_path}\n"
        f"\n"
        f"nc: {len(YOLO_CLASS_NAMES)}\n"
        f"names: {YOLO_CLASS_NAMES}\n"
    )
    yaml_path.write_text(content)
    logger.info("Wrote dataset YAML: %s", yaml_path)
    return yaml_path


def train() -> Dict[str, Any]:
    """Run full training pipeline: convert → train → export.

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

    stats: Dict[str, Any] = {}

    # ── Step 0: Clean stale YOLO dirs ─────────────────────────────────────────
    _clean_yolo_dirs()

    # ── Step 1: Convert COCO → YOLO ──────────────────────────────────────────
    logger.info("Converting COCO annotations to YOLO format...")
    for split in ("train", "val"):
        n = _convert_coco_to_yolo(split)
        stats[f"{split}_images"] = n
        logger.info("  %s: %d images (incl. oversampling)", split, n)

    # ── Step 2: Write dataset.yaml ───────────────────────────────────────────
    yaml_path = _write_dataset_yaml()

    # ── Step 3: Train ────────────────────────────────────────────────────────
    device = _detect_device()
    logger.info("Starting YOLO12 training on device=%s ...", device)
    model = YOLO(YOLO_BASE_MODEL)

    results = model.train(
        data=str(yaml_path),
        epochs=YOLO_EPOCHS,
        imgsz=YOLO_IMGSZ,
        batch=YOLO_BATCH,
        device=device,
        patience=YOLO_PATIENCE,
        cos_lr=True,
        cls=YOLO_CLS_LOSS,
        save_period=YOLO_SAVE_PERIOD,
        project=str(YOLO_DIR / "runs"),
        name="train",
        exist_ok=True,
        verbose=True,
    )

    # Find best weights
    train_dir = Path(results.save_dir)
    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)

    # ── Step 4: Export to ONNX ───────────────────────────────────────────────
    logger.info("Exporting to ONNX...")
    best_model = YOLO(str(best_pt))
    export_path = best_model.export(
        format="onnx",
        imgsz=YOLO_IMGSZ,
        simplify=True,
        opset=ONNX_OPSET,
    )
    logger.info("ONNX exported: %s", export_path)

    # Validate ONNX output shape
    _validate_onnx(export_path)

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

    # ── Step 6: Read training metrics if available ───────────────────────────
    results_csv = train_dir / "results.csv"
    if results_csv.exists():
        _log_metrics(results_csv, stats)

    return stats


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)

    # Expect input (1, 3, H, W) and output (1, 4+num_classes, N)
    if len(out.shape) != 3:
        raise ValueError(
            f"Unexpected ONNX output rank {len(out.shape)}: {out.shape}. "
            f"Expected (1, 4+C, N)."
        )

    num_classes = out.shape[1] - 4 if isinstance(out.shape[1], int) else None
    if num_classes is not None and num_classes != len(YOLO_CLASS_NAMES):
        raise ValueError(
            f"ONNX output has {num_classes} classes but config expects "
            f"{len(YOLO_CLASS_NAMES)}: {YOLO_CLASS_NAMES}"
        )

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


def _log_metrics(results_csv: Path, stats: Dict[str, Any]):
    """Read final metrics from results.csv and add to stats."""
    import csv

    with open(results_csv) as f:
        reader = csv.DictReader(f)
        rows = list(reader)

    if not rows:
        return

    last = rows[-1]
    # Ultralytics CSV headers have leading spaces
    for key in last:
        clean = key.strip()
        if "mAP50" in clean and "mAP50-95" not in clean:
            try:
                stats["mAP50"] = float(last[key])
            except (ValueError, TypeError):
                pass
        elif "mAP50-95" in clean:
            try:
                stats["mAP50_95"] = float(last[key])
            except (ValueError, TypeError):
                pass

    if "mAP50" in stats:
        logger.info("Final mAP50: %.4f", stats["mAP50"])
    if "mAP50_95" in stats:
        logger.info("Final mAP50-95: %.4f", stats["mAP50_95"])


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    train()
