# scripts/train_yolo_seg.py
"""
Обучение YOLOv8 на сегментацию.
Принимает параметры через командную строку.
"""

import sys
import time
import argparse
from pathlib import Path

import torch
from ultralytics import YOLO


def check_environment():
    print("=" * 60)
    print("ПРОВЕРКА ОКРУЖЕНИЯ (SEGMENT)")
    print("=" * 60)
    print(f"Python: {sys.version.split()[0]}")
    print(f"PyTorch: {torch.__version__}")
    print(f"CUDA доступна: {torch.cuda.is_available()}")
    if 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:.2f} GB")
    print("=" * 60)


def train(data_path: str, model: str, epochs: int, imgsz: int,
          batch: int, device: str, project: str, name: str):
    check_environment()

    if '-seg' not in model.lower():
        print(f"⚠️  Модель '{model}' не является сегментационной!")

    log_dir = Path(project) / name
    log_dir.mkdir(parents=True, exist_ok=True)

    print(f"\nЗагрузка модели: {model}")
    yolo = YOLO(model)

    args = {
        'task': 'segment',
        'data': data_path,
        'epochs': epochs,
        'imgsz': imgsz,
        'batch': batch,
        'device': device,
        'project': project,
        'name': name,
        'exist_ok': True,
        'pretrained': True,
        'optimizer': 'auto',
        'lr0': 0.01,
        'lrf': 0.01,
        'momentum': 0.937,
        'weight_decay': 0.0005,
        'warmup_epochs': 3.0,
        'warmup_momentum': 0.8,
        'warmup_bias_lr': 0.1,
        'box': 7.5,
        'cls': 0.5,
        'dfl': 1.5,
        'mask_ratio': 4,
        'overlap_mask': True,
        'patience': 20,
        'save': True,
        'save_period': 10,
        'plots': True,
        'verbose': True,
    }

    print(f"\nЗапуск обучения СЕГМЕНТАЦИИ...")
    print(f"Датасет: {data_path}")
    print(f"Модель: {model}")
    print(f"Эпохи: {epochs}, Batch: {batch}, imgsz: {imgsz}")
    print(f"Результаты: {log_dir}")
    print("=" * 60)

    t0 = time.time()
    results = yolo.train(**args)
    print(f"\nВремя обучения: {(time.time() - t0) / 60:.1f} мин")
    print(f"Лучшая модель: {log_dir / 'weights' / 'best.pt'}")
    return results


def main():
    parser = argparse.ArgumentParser(description='Обучение YOLOv8 segment')
    parser.add_argument('--data', type=str, required=True, help='Путь к data.yaml')
    parser.add_argument('--model', type=str, default='yolov8s-seg.pt')
    parser.add_argument('--epochs', type=int, default=100)
    parser.add_argument('--imgsz', type=int, default=512)
    parser.add_argument('--batch', type=int, default=16)
    parser.add_argument('--device', type=str, default='0')
    parser.add_argument('--project', type=str, default='runs/segment')
    parser.add_argument('--name', type=str, default='seg_experiment')
    args = parser.parse_args()

    train(args.data, args.model, args.epochs, args.imgsz,
          args.batch, args.device, args.project, args.name)


if __name__ == '__main__':
    main()