# ============================================================
# CELL 1 — Install & verify GPU
# ============================================================
import torch
print(f"PyTorch   : {torch.__version__}")
print(f"CUDA      : {torch.cuda.is_available()}")
print(f"GPU       : {torch.cuda.get_device_name(0)}")
print(f"VRAM      : {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")


# ============================================================
# CELL 2 — Fix data.yaml paths for Kaggle  ← FIXED
# ============================================================

import os
import yaml

# ✅ Your exact dataset path (from your screenshot)
DATASET_DIR = "merged_dataset"
DATA_YAML   = os.path.join(DATASET_DIR, "data.yaml")

print("Verifying dataset structure:")
for item in os.listdir(DATASET_DIR):
    full_path = os.path.join(DATASET_DIR, item)
    if os.path.isdir(full_path):
        print(f"  [folder] {full_path}")
    else:
        print(f"  [file]   {full_path}")

# ── Read existing yaml ────────────────────────────────────────
with open(DATA_YAML, "r") as f:
    cfg = yaml.safe_load(f)

print(f"\nOriginal yaml paths:")
print(f"  train : {cfg.get('train')}")
print(f"  val   : {cfg.get('val')}")
print(f"  test  : {cfg.get('test')}")

# ── Fix paths to absolute Kaggle paths ───────────────────────
# Kaggle cannot resolve relative paths like ../train/images
cfg["train"] = f"{DATASET_DIR}/train/images"
cfg["val"]   = f"{DATASET_DIR}/valid/images"
cfg["test"]  = f"{DATASET_DIR}/test/images"

# ── Save fixed yaml to /kaggle/working (writable folder) ─────
FIXED_YAML = "/kaggle/working/data_fixed.yaml"
with open(FIXED_YAML, "w") as f:
    yaml.dump(cfg, f, default_flow_style=False)

print(f"\nFixed yaml saved to: {FIXED_YAML}")
print(f"  train : {cfg['train']}")
print(f"  val   : {cfg['val']}")
print(f"  test  : {cfg['test']}")
print(f"  nc    : {cfg.get('nc')}")
print(f"  names : {cfg.get('names')}")

# ── Verify image folders actually exist ──────────────────────
for split, path in [("train", cfg["train"]), ("val", cfg["val"]), ("test", cfg["test"])]:
    exists = os.path.isdir(path)
    count  = len(os.listdir(path)) if exists else 0
    print(f"  {split:5s} images: {'OK' if exists else 'MISSING'} ({count} files)")


# ============================================================
# CELL 3 — Train on GPU
# ============================================================

from ultralytics import YOLO
from pathlib import Path

# ✅ Use the fixed yaml (absolute paths)
DATA_YAML_FIXED = "data.yaml"
OUTPUT_DIR      = "result_directorygit"

# ── Model ────────────────────────────────────────────────────
# T4 has 15.6 GB VRAM — yolov8m is the best quality/speed balance
MODEL_SIZE = "yolo26n.pt"
model = YOLO(MODEL_SIZE)

# ── Hyperparameters — tuned for Kaggle T4 GPU ────────────────
results = model.train(

    # ── Data ──────────────────────────────────────────────────
    data        = DATA_YAML_FIXED,
    imgsz       = 640,          # full 640 on GPU
    rect        = False,

    # ── Training schedule ──────────────────────────────────────
    epochs      = 100,          # ~15 min on T4
    patience    = 20,
    batch       = 32,           # T4 can handle batch=32 easily
    workers     = 4,

    # ── Optimizer ─────────────────────────────────────────────
    optimizer   = "AdamW",
    lr0         = 0.001,
    lrf         = 0.01,
    momentum    = 0.937,
    weight_decay= 0.0005,

    # ── LR Warmup ─────────────────────────────────────────────
    warmup_epochs   = 3.0,
    warmup_momentum = 0.8,
    warmup_bias_lr  = 0.1,

    # ── Loss weights ──────────────────────────────────────────
    box         = 7.5,
    cls         = 0.5,
    dfl         = 1.5,

    # ── Augmentation — full GPU augmentation ──────────────────
    hsv_h       = 0.015,
    hsv_s       = 0.7,
    hsv_v       = 0.4,
    degrees     = 10.0,
    translate   = 0.1,
    scale       = 0.5,
    shear       = 2.0,
    perspective = 0.0,
    flipud      = 0.1,
    fliplr      = 0.5,
    mosaic      = 1.0,
    mixup       = 0.15,
    copy_paste  = 0.1,

    # ── Output ────────────────────────────────────────────────
    project     = OUTPUT_DIR,
    name        = "yolov8_microplastic",
    exist_ok    = True,
    save        = True,
    save_period = 20,
    plots       = True,
    verbose     = True,
    cache       = True,         # fast SSD on Kaggle — safe to cache

    # ── Device ────────────────────────────────────────────────
    device      = 0,            # GPU 0 (Tesla T4)
    amp         = True,         # Mixed precision — faster on GPU
    multi_scale = False,
)

print("\nTraining complete!")


# ============================================================
# CELL 4 — Evaluate on test set
# ============================================================

best_weights = Path(OUTPUT_DIR) / "yolov8_microplastic" / "weights" / "best.pt"
best_model   = YOLO(str(best_weights))

test_metrics = best_model.val(
    data      = DATA_YAML_FIXED,
    split     = "test",
    imgsz     = 640,
    batch     = 32,
    conf      = 0.25,
    iou       = 0.6,
    device    = 0,
    plots     = True,
    save_json = True,
)

print("\n── TEST RESULTS ──────────────────────────────")
print(f"  mAP50     : {test_metrics.box.map50:.4f}")
print(f"  mAP50-95  : {test_metrics.box.map:.4f}")
print(f"  Precision : {test_metrics.box.mp:.4f}")
print(f"  Recall    : {test_metrics.box.mr:.4f}")
print("──────────────────────────────────────────────")


# ============================================================
# CELL 5 — List output files for download
# ============================================================

print("\nFiles ready to download from Output panel:")
for f in sorted(Path(OUTPUT_DIR).rglob("*")):
    if f.suffix in [".pt", ".yaml", ".png", ".csv", ".json"]:
        size_mb = f.stat().st_size / 1e6
        print(f"  {str(f).replace(OUTPUT_DIR, '')}  ({size_mb:.1f} MB)")

print("\nTo download best.pt:")
print("  Output panel (right side) → yolov8_microplastic → weights → best.pt → Download")