from pathlib import Path

from ultralytics import YOLO


DATASET_YAML = Path(__file__).parent / "dataset" / "dataset.yaml"
RUNS_DIR = Path(__file__).parent / "runs"


def train():
    model = YOLO("runs/pothole/weights/best.pt")

    results = model.train(
        data=str(DATASET_YAML),
        epochs=100,
        imgsz=1280,
        batch=64,
        project=str(RUNS_DIR),
        name="pothole",
        exist_ok=True,

        # LR scheduling — cosine annealing with warmup
        lr0=0.001,
        lrf=0.01,
        warmup_epochs=3,
        warmup_momentum=0.8,
        cos_lr=True,

        # Regularization
        weight_decay=0.0005,
        dropout=0.0,
        patience=50,

        # Augmentation — tuned for road/pothole data
        mosaic=1.0,
        copy_paste=0.3,       # paste potholes onto new backgrounds
        degrees=5.0,           # slight rotation (roads tilt a bit)
        scale=0.5,             # scale jitter
        fliplr=0.5,            # horizontal flip
        flipud=0.0,            # roads are always right-side up
        hsv_h=0.015,           # hue jitter (lighting conditions)
        hsv_s=0.7,             # saturation jitter (wet/dry roads)
        hsv_v=0.4,             # brightness jitter (day/night)
        translate=0.1,
        perspective=0.0,       # road camera is always flat
        mixup=0.1,             # blend images for harder examples
    )

    print(f"Training complete. Results saved to: {results.save_dir}")
    return results


def main():
    train()


if __name__ == "__main__":
    main()