"""
STEP 3 - train.py
================

Fine-tunes YOLOv8-nano on your augmented dataset.
Pre-trained COCO weights give you a huve head-start - the backbone
already knows edges, textures, and shapes; it only needs to learn
what Famico (and other products) look like.

Install:
    pip install ultralytics

Usage:
    python train.py

Outputs:
    runs/detect/coffee_v1/weights/best.pt <- use this for detection
"""

import os
import torch
from pathlib import Path
from ultralytics import YOLO

# ── Configuration ───────────────────────────────────────────────────
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

DATASET_YAML = "dataset.yml" # Generated by prepare_dataset.py
MODEL        = "yolov8n.pt"  # nano - fast, good for few-shot
                             # alternativse to try: yolov8s.pt (better acc)
                             #                      yolov8m.pt (even better, needs GPU)

PROJECT      = "runs/detect"
RUN_NAME     = "coffee_v1"
IMG_SIZE     = 640


# ── Training hyperparameters ───────────────────────────────────────────────────
#
# Tuned for: small dataset (few-shot), packaged goods, retail shelf.
#
EPOCHS        = 150
PATIENCE      = 30        # Early Stopping - stops if no improvement
BATCH         = 8         # Lower if you get OOM on CPU/small GPU.
LR0           = 0.005     # Initial LR (Learing Rate) - lower than default (0.01)

LRF           = 0.01      # final LR = LR0 * LRF
MOMENTUM      = 0.938
WEIGHT_DECAY  = 0.0005
WARMUP_EPOCHS = 5         # Longer warmup stabilises few-shot training.
IOU_THRESH    = 0.5       # IoU threshold for NMS.
CONF_THRESH   = 0.25      # confidence threshold at val time.


# ── YOLOv8 built-in augmentation (on top of our albumentatios pass) ────────────
#
# These run during tarining batches - complementary to augment.py.
# Mosaic is the single most impactful augmentation for small datasets.
#
MOSAIC     = 1.0    # mosaic probability
MIXUP      = 0.1    # mixup probability
COPY_PASTE = 0.1
DEGREES    = 5.0    # degrees ±5°
TRANSLATE  = 0.1
SCALE      = 0.5
FLIPUD     = 0.0
FLIPLR     = 0.5
HSV_H      = 0.015
HSV_S      = 0.7
HSV_V      = 0.4

def check_enviroment():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cpu":
        print("[INFO] No GPU detected - trainning on CPU")
        print("       Expected time: ~60-90min for 150 epochs on a typical laptop.")
        print("       Tips: use Google Colab or Kaggle they got free GPU, if training is too slow.\n")
    else:
        gpu= torch.cuda.get_device_name(0)
        print(f"[INFO] GPU detected: {gpu}")
        print("        Expected time: ~5-15min.\n")
    return device

def train():
    device = check_enviroment()

    if not Path(DATASET_YAML).exists():
        print(f"[ERROR] {DATASET_YAML} not found.")
        print("         Run prepare_dataset.py first.")
        return

    model = YOLO(MODEL)
    print(f"Starting training: {EPOCHS} epochs, img size {IMG_SIZE}")
    print(f"Dataset: {DATASET_YAML}")
    print(f"Output:  {PROJECT}/{RUN_NAME}/weights/best.pt\n")

    results = model.train(
            data          = DATASET_YAML,
            epochs        = EPOCHS,
            patience      = PATIENCE,
            imgsz         = IMG_SIZE,
            batch         = BATCH,
            device        = device,
            project       = PROJECT,
            name          = RUN_NAME,
            exist_ok      = True,

            lr0           = LR0,
            lrf           = LRF,
            momentum      = MOMENTUM,
            weight_decay  = WEIGHT_DECAY,
            warmup_epochs = WARMUP_EPOCHS,

            iou           = IOU_THRESH,
            conf          = CONF_THRESH,

            # Built-in augmentation
            mosaic        = MOSAIC,
            mixup         = MIXUP,
            copy_paste    = COPY_PASTE,
            degrees       = DEGREES,
            translate     = TRANSLATE,
            scale         = SCALE,
            flipud        = FLIPUD,
            fliplr        = FLIPLR,
            hsv_h         = HSV_H,
            hsv_s         = HSV_S,
            hsv_v         = HSV_V,

            save          = True,
            save_period   = 10,    # Checkpoint every 10 epochs
            plots         = True,  # Saves training curvse as PNGs
            verbose       = True,
            seed          = 42,
    )

    best_weights = Path(PROJECT) / RUN_NAME / "weights" / "best.pt"
    print("\n" + "-" * 60)
    print(f"Training complete.")
    print(f"Best weights saved -> {best_weights}")
    print(f"mAP50: {results.results_dict.get('metrics/mAP50(B)', 'N/A'):.3f}")
    print(f"mAP50-95: {results.results_dict.get('metrics/mAP50-95(B)', 'N/A'):.3f}")
    print("-" * 60)
    print("\nRun next: python detect.py --source scene.jpg")

    return best_weights

if __name__ == "__main__":
    train()