"""
train_from_checkpoint.py
========================
Resume training from last saved checkpoint runs/detect/coffee_v1/weights/last.pt

Usage:
    python train_from_checkpoint.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

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 fix_grad_scaler(checkpoint_path: str):
    """
    Patches a checkpoint whose GradScaler state is empty
    (happens when last.pt was saved on CPU / AMP-disabled env).
    """
    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)

    default_scaler_state = {
        "scale":        torch.tensor(65536.0),   # standard initial scale
        "growth_factor": 2.0,
        "backoff_factor": 0.5,
        "growth_interval": 2000,
        "_growth_tracker": 0,
    }

    patched = False
    # Ultralytics stores the scaler under the key 'scaler'
    if "scaler" in ckpt:
        if not ckpt["scaler"]:          # empty dict → broken
            ckpt["scaler"] = default_scaler_state
            patched = True
    else:
        ckpt["scaler"] = default_scaler_state
        patched = True

    if patched:
        torch.save(ckpt, checkpoint_path)
        print(f"[INFO] GradScaler state patched in {checkpoint_path}")
    else:
        print(f"[INFO] GradScaler state looks fine, no patch needed.")

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_path = Path(PROJECT) / Path(PROJECT) / RUN_NAME / "weights" / "last.pt"
    if not model_path.exists():
        print(f"[ERROR] {model_path} not found.")
        print("         Run train.py first.")
        return

    # Patch checkpoint before loading.
    if device == "cuda":
        fix_grad_scaler(str(model_path))

    model = YOLO(model_path)
    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(
            resume        = True,
            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()