#!/usr/bin/env python3
"""
YOLO Fine-tuning для Subnet 44 ScoreVision
Запускать на clore.ai (RTX 5090)
Вывод: runs/detect/train/weights/best.pt
"""

import os
from ultralytics import YOLO
from huggingface_hub import snapshot_download

DATASET_DIR = "/training/soccernet_v3_h250"
YAML_PATH = os.path.join(DATASET_DIR, "dataset.yaml")
EPOCHS = 100
IMAGE_SIZE = 640
BATCH_SIZE = 32

# ========== Шаг 1: Скачиваем датасет ==========
print("⬇️  Скачиваем датасет SoccerNet...")
if not os.path.exists(os.path.join(DATASET_DIR, "images")):
    snapshot_download(
        repo_id="SoccerNet/SoccerNet",
        repo_type="dataset",
        local_dir=DATASET_DIR,
        allow_patterns=["images/**", "labels/**"],
    )
print("✅ Датасет готов")

# ========== Шаг 2: dataset.yaml ==========
if not os.path.exists(YAML_PATH):
    import yaml
    # Определяем классы из лейблов
    names = {}
    for split in ["train", "val"]:
        dirs = [os.path.join(DATASET_DIR, "labels", split)]
        for d in dirs:
            if os.path.isdir(d):
                for f in os.listdir(d)[:1]:
                    with open(os.path.join(d, f)) as fh:
                        for line in fh:
                            p = line.strip().split()
                            if p:
                                names[int(p[0])] = f"class_{p[0]}"
                break
    if not names:
        names = {0: "ball", 1: "player"}

    with open(YAML_PATH, "w") as f:
        yaml.dump({
            "path": DATASET_DIR,
            "train": "images/train",
            "val": "images/val",
            "names": names
        }, f)
    print(f"✅ Создан dataset.yaml: классы = {names}")

# ========== Шаг 3: Файн-тюнинг ==========
print("🔧 Загружаем YOLO11s...")
model = YOLO("yolo11s.pt")

print(f"🏋️  Файн-тюнинг: {EPOCHS} эпох, размер {IMAGE_SIZE}, batch {BATCH_SIZE}...")
model.train(
    data=YAML_PATH,
    epochs=100,
    imgsz=640,
    batch=32,
    device=0,
    cache=True,
    augment=True,
)
print("✅ Обучение завершено!")

# ========== Шаг 4: Оценка ==========
metrics = YOLO("runs/detect/train/weights/best.pt").val(data=YAML_PATH, device=0)
print(f"\n📊 mAP50: {metrics.box.map50:.4f} | mAP50-95: {metrics.box.map:.4f}")

# ========== Шаг 5: Информация для скачивания ==========
import json
from pathlib import Path

weights_dir = Path("runs/detect/train/weights")
files = {f.name: f.stat().st_size for f in weights_dir.glob("*.pt")}
print(f"\n📦 Файлы модели в: {weights_dir.absolute()}")
print(f"\n📋 СКАЧАТЬ:\nscp -r {weights_dir.parent.absolute()} user@server:/tmp/")