"""
Fine-tune a YOLO26n object-detection model on a custom dataset.

Expected dataset configuration:
    data.yaml

Example data.yaml:
    path: /absolute/path/to/dataset
    train: images/train
    val: images/val

    names:
      0: person
      1: car
      2: bicycle

The corresponding label files should use standard YOLO detection format:
    class_id x_center y_center width height

All coordinates must be normalized to the range 0.0-1.0.
"""

from pathlib import Path

import torch
from ultralytics import YOLO


# ============================================================================
# User-configurable settings
# ============================================================================

# Path to the pretrained model checkpoint.
# If the file is not found locally, Ultralytics may download it automatically.
MODEL_PATH = "yolo26n.pt"

# Path to the dataset
DATASET_PATH = Path("C:/Teaching/ComputerVision/examples/yolo-custom/bigdata")

# Path to your dataset configuration file.
DATA_CONFIG_PATH = DATASET_PATH / "data.yaml"

# Number of complete passes through the training dataset.
EPOCHS = 50

# Input image size in pixels.
# Common values are 640, 768, or 1024.
IMAGE_SIZE = 320
# IMAGE_SIZE = 640

# Number of images processed at once.
# Reduce this if you run out of GPU memory.
# Set this to -1 to let Ultralytics automatically select a batch size.
BATCH_SIZE = 16

# Training device:
#   "cpu"  -> CPU
#   0      -> first CUDA GPU
#   1      -> second CUDA GPU
#   "mps"  -> Apple Silicon GPU
#   [0, 1] -> multiple CUDA GPUs
DEVICE = 0 if torch.cuda.is_available() else "cpu"

# Directory where training results will be saved.
PROJECT_DIR = DATASET_PATH / "runs"

# Name of this particular training run.
RUN_NAME = f"yolo26n_{DATASET_PATH.name}"

# Number of epochs without validation improvement before early stopping.
PATIENCE = 30

# Save the best and latest checkpoints during training.
SAVE_CHECKPOINTS = True

# Whether to overwrite an existing run directory with the same name.
OVERWRITE_EXISTING_RUN = False

# Number of dataloader worker processes.
# Use 0 on Windows if you encounter multiprocessing problems.
WORKERS = 8

# Random seed for more reproducible experiments.
SEED = 42

# Use pretrained weights from MODEL_PATH.
# This should remain True for fine-tuning.
USE_PRETRAINED_WEIGHTS = True

# Optimizer selection.
# "auto" lets Ultralytics select an optimizer.
# Other possible choices may include "SGD", "Adam", or "AdamW".
OPTIMIZER = "auto"

# Initial learning rate.
# Set to None to let Ultralytics choose it automatically.
LEARNING_RATE = None

# Final learning-rate multiplier relative to the initial learning rate.
LEARNING_RATE_FINAL_MULTIPLIER = 0.01

# Weight decay helps regularize the model.
WEIGHT_DECAY = 0.0005

# Enable mixed-precision training when supported by the hardware.
# This usually reduces GPU memory usage and can improve training speed.
USE_AMP = True

# Cache images in memory or on disk to speed up training.
# Options:
#   False -> no caching
#   True  -> cache in RAM
#   "disk" -> cache on disk
CACHE_IMAGES = False

# Use rectangular batches. This can reduce padding for images with
# substantially different aspect ratios.
RECTANGULAR_BATCHES = False

# Number of images used for validation in each batch.
VALIDATION_BATCH_SIZE = None  # None means use the training batch size.

# Close mosaic augmentation during the final N epochs.
CLOSE_MOSAIC_EPOCHS = 10

# Data augmentation settings.
# Adjust these for your specific image domain.
MOSAIC_PROBABILITY = 1.0
MIXUP_PROBABILITY = 0.0
COPY_PASTE_PROBABILITY = 0.0
HORIZONTAL_FLIP_PROBABILITY = 0.5
VERTICAL_FLIP_PROBABILITY = 0.0
ROTATION_DEGREES = 0.0
TRANSLATION_FRACTION = 0.1
SCALE_RANGE = 0.5

# Optional: freeze the first N model layers.
# Set to 0 or None to train all layers.
# Freezing can help with very small datasets.
# FREEZE_LAYERS = None
FREEZE_LAYERS = 10

# Automatically validate the best model after training.
RUN_FINAL_VALIDATION = True

# Optional path for exporting the trained model.
# Set to None to skip exporting.
# Examples:
# EXPORT_FORMAT = None
EXPORT_FORMAT = "onnx"
# EXPORT_FORMAT = "torchscript"
# EXPORT_FORMAT = "engine"


# ============================================================================
# Helper functions
# ============================================================================

def choose_device():
    """
    Select and display the available training device.

    The configured DEVICE value takes precedence, except when it is a CUDA
    device that is unavailable. In that case, training falls back to CPU.
    """

    # If the user explicitly requested MPS, verify that MPS is available.
    if DEVICE == "mps":
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"

        print("MPS is unavailable. Falling back to CPU.")
        return "cpu"

    # If the user explicitly requested CUDA, verify CUDA availability.
    if isinstance(DEVICE, int):
        if torch.cuda.is_available():
            return DEVICE

        print("CUDA is unavailable. Falling back to CPU.")
        return "cpu"

    # For strings such as "cpu" or "0,1", use the configured value directly.
    return DEVICE


def build_training_arguments(device):
    """
    Build the keyword arguments passed to model.train().

    Keeping the settings in a dictionary makes it easy to add, remove, or
    modify training options.
    """

    # Required training arguments.
    train_args = {
        "data": str(DATA_CONFIG_PATH),
        "epochs": EPOCHS,
        "imgsz": IMAGE_SIZE,
        "batch": BATCH_SIZE,
        "device": device,
        "project": str(PROJECT_DIR),
        "name": RUN_NAME,
        "patience": PATIENCE,
        "save": SAVE_CHECKPOINTS,
        "exist_ok": OVERWRITE_EXISTING_RUN,
        "workers": WORKERS,
        "seed": SEED,
        "optimizer": OPTIMIZER,
        "weight_decay": WEIGHT_DECAY,
        "amp": USE_AMP,
        "cache": CACHE_IMAGES,
        "rect": RECTANGULAR_BATCHES,
        "close_mosaic": CLOSE_MOSAIC_EPOCHS,
        "mosaic": MOSAIC_PROBABILITY,
        "mixup": MIXUP_PROBABILITY,
        "copy_paste": COPY_PASTE_PROBABILITY,
        "fliplr": HORIZONTAL_FLIP_PROBABILITY,
        "flipud": VERTICAL_FLIP_PROBABILITY,
        "degrees": ROTATION_DEGREES,
        "translate": TRANSLATION_FRACTION,
        "scale": SCALE_RANGE,
    }

    # Add the validation batch size only when explicitly configured.
    if VALIDATION_BATCH_SIZE is not None:
        train_args["val_batch"] = VALIDATION_BATCH_SIZE

    # Add the learning-rate settings only when explicitly configured.
    if LEARNING_RATE is not None:
        train_args["lr0"] = LEARNING_RATE
        train_args["lrf"] = LEARNING_RATE_FINAL_MULTIPLIER

    # Add layer freezing only when explicitly configured.
    if FREEZE_LAYERS is not None and FREEZE_LAYERS > 0:
        train_args["freeze"] = FREEZE_LAYERS

    return train_args


# ============================================================================
# Main training procedure
# ============================================================================

def main():
    """
    Load the pretrained checkpoint, fine-tune it, validate it, and optionally
    export the trained model.
    """

    # Convert paths to Path objects for simple existence checks.
    model_path = Path(MODEL_PATH)
    data_path = Path(DATA_CONFIG_PATH)

    # Confirm that the dataset YAML exists before starting.
    if not data_path.exists():
        raise FileNotFoundError(
            f"Dataset YAML was not found: {data_path.resolve()}"
        )

    # If the model checkpoint is not present locally, Ultralytics can often
    # download official pretrained weights when YOLO() is called.
    if not model_path.exists():
        print(
            f"Model file '{MODEL_PATH}' was not found locally. "
            "Ultralytics will try to download it."
        )

    # Select the training device.
    device = choose_device()

    # Print useful information before training begins.
    print(f"Using device: {device}")
    print(f"Model: {MODEL_PATH}")
    print(f"Dataset configuration: {DATA_CONFIG_PATH}")
    print(f"Epochs: {EPOCHS}")
    print(f"Image size: {IMAGE_SIZE}")
    print(f"Batch size: {BATCH_SIZE}")

    # Load the YOLO26n pretrained model.
    # This is the PyTorch/Ultralytics model interface.
    model = YOLO(MODEL_PATH)

    # Fine-tune the model on your custom dataset.
    train_args = build_training_arguments(device)

    # Start training.
    training_results = model.train(**train_args)

    # Use the directory that Ultralytics actually selected.
    run_directory = Path(model.trainer.save_dir)
    best_model_path = run_directory / "weights" / "best.pt"

    # Print the location of the run directory.
    print(f"\nTraining finished. Results were saved to: {run_directory}")

    # The best checkpoint is normally saved here.
    best_model_path = run_directory / "weights" / "best.pt"

    # Validate the best checkpoint after training.
    if RUN_FINAL_VALIDATION and best_model_path.exists():
        print(f"\nValidating best checkpoint: {best_model_path}")

        # Load the best-performing checkpoint rather than the final epoch.
        best_model = YOLO(str(best_model_path))

        # Run validation using the same dataset YAML.
        validation_results = best_model.val(
            data=DATA_CONFIG_PATH,
            imgsz=IMAGE_SIZE,
            batch=(
                VALIDATION_BATCH_SIZE
                if VALIDATION_BATCH_SIZE is not None
                else BATCH_SIZE
            ),
            device=device,
            workers=WORKERS,
            project=str(run_directory),
        )

        # Display the main validation metric when available.
        if hasattr(validation_results, "box"):
            print(f"mAP50-95: {validation_results.box.map:.4f}")
            print(f"mAP50:    {validation_results.box.map50:.4f}")

    # Optionally export the best checkpoint to another deployment format.
    if EXPORT_FORMAT is not None and best_model_path.exists():
        print(f"\nExporting best checkpoint to {EXPORT_FORMAT} format.")

        export_model = YOLO(str(best_model_path))
        exported_path = export_model.export(
            format=EXPORT_FORMAT,
            imgsz=IMAGE_SIZE,
            device=device,
        )

        print(f"Exported model: {exported_path}")

    # Keep the returned training object available if you want to inspect
    # additional metrics programmatically.
    return training_results


# This prevents the training code from running automatically if this file
# is imported as a module from another Python script.
if __name__ == "__main__":
    main()