# scripts/ttv_freq2.py
"""
Версия 6: Читает predefined_classes.txt для сохранения полной схемы классов.
Это позволяет использовать общую схему 12 классов для всех экспериментов,
даже если в конкретном датасете присутствуют не все классы.
"""

import os
import shutil
import random
import argparse
from pathlib import Path

try:
    import yaml
except ImportError:
    raise SystemExit("Установите PyYAML: pip install pyyaml")


def load_predefined_classes(path: Path) -> dict:
    """
    Читает predefined_classes.txt и возвращает маппинг:
    {имя_класса: (id_1based, yolo_index_0based)}
    
    Формат файла (одна строка = один класс):
        right_kidney
        right_kidney_upper_segment
        ...
    """
    class_map = {}
    with open(path, 'r', encoding='utf-8') as f:
        for idx, line in enumerate(f):
            name = line.strip()
            if name:
                class_map[name] = {
                    'id_1based': idx + 1,
                    'yolo_idx': idx,
                }
    return class_map


def clean_caches(output_dir: Path):
    """Удаляет старые .cache файлы."""
    for cache in output_dir.glob('*.cache'):
        try:
            cache.unlink()
            print(f"  Удалён кэш: {cache.name}")
        except OSError:
            pass


def remap_labels(src_dir: Path, dst_dir: Path, id_to_yolo_idx: dict):
    """
    Копирует .txt файлы с перенумерацией классов из 1-based в 0-based.
    """
    if dst_dir.exists():
        shutil.rmtree(dst_dir)
    dst_dir.mkdir(parents=True, exist_ok=True)

    copied = 0
    skipped = 0
    for txt_file in src_dir.glob('*.txt'):
        with open(txt_file, 'r', encoding='utf-8') as f:
            lines = f.readlines()

        new_lines = []
        for line in lines:
            parts = line.strip().split()
            if not parts:
                continue
            try:
                old_class = int(parts[0])
                if old_class in id_to_yolo_idx:
                    parts[0] = str(id_to_yolo_idx[old_class])
                    new_lines.append(' '.join(parts))
                else:
                    skipped += 1
            except ValueError:
                continue

        if new_lines:
            with open(dst_dir / txt_file.name, 'w', encoding='utf-8') as f:
                f.write('\n'.join(new_lines) + '\n')
            copied += 1

    return copied, skipped


def create_split(
    images_dir: Path,
    labels_dir: Path,
    output_dir: Path,
    task: str,
    predefined_classes: dict,
    train_ratio: float = 0.8,
    val_ratio: float = 0.1,
    seed: int = 42,
):
    random.seed(seed)

    images_dir = images_dir.resolve()
    labels_dir = labels_dir.resolve()
    output_dir = output_dir.resolve()

    if not images_dir.exists():
        raise FileNotFoundError(f"Папка с изображениями не найдена: {images_dir}")
    if not labels_dir.exists():
        raise FileNotFoundError(f"Папка с лейблами не найдена: {labels_dir}")

    print("\nОчистка кэша...")
    clean_caches(output_dir)

    image_files = [f for f in os.listdir(images_dir)
                   if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    if not image_files:
        raise ValueError(f"В {images_dir} не найдено изображений.")

    print(f"\n{'='*60}")
    print(f"Задача: {task.upper()}")
    print(f"{'='*60}")
    print(f"Изображения: {images_dir}")
    print(f"Лейблы: {labels_dir}")
    print(f"Найдено изображений: {len(image_files)}")

    # Собираем статистику по классам
    class_to_images = {}
    classes_found = set()

    for image_file in image_files:
        label_file = os.path.splitext(image_file)[0] + '.txt'
        label_path = labels_dir / label_file

        if label_path.exists():
            with open(label_path, 'r', encoding='utf-8') as f:
                for line in f:
                    parts = line.strip().split()
                    if not parts:
                        continue
                    try:
                        class_label = int(parts[0])
                        classes_found.add(class_label)
                        class_to_images.setdefault(class_label, set()).add(image_file)
                    except ValueError:
                        continue

    # Маппинг 1-based ID -> 0-based YOLO index
    id_to_yolo_idx = {info['id_1based']: info['yolo_idx']
                      for info in predefined_classes.values()}

    print(f"\n📋 Схема классов (из predefined_classes.txt):")
    print(f"  Всего классов: {len(predefined_classes)}")
    present = []
    absent = []
    for name, info in predefined_classes.items():
        if info['id_1based'] in classes_found:
            present.append((name, info, len(class_to_images[info['id_1based']])))
        else:
            absent.append((name, info))

    print(f"\n  ✓ Присутствуют в датасете ({len(present)}):")
    for name, info, count in present:
        print(f"    [{info['id_1based']:2d}] → YOLO idx {info['yolo_idx']:2d} : {name:<40s} ({count} изобр.)")

    if absent:
        print(f"\n  ○ Отсутствуют ({len(absent)}):")
        for name, info in absent:
            print(f"    [{info['id_1based']:2d}] → YOLO idx {info['yolo_idx']:2d} : {name}")

    # Разбиение по наборам (стратифицированное)
    sorted_classes = sorted(class_to_images, key=lambda k: len(class_to_images[k]))

    train_files, val_files, test_files = set(), set(), set()

    print(f"\n--- Распределение по наборам ---")
    for class_label in sorted_classes:
        images = set(class_to_images[class_label])
        total = len(images)

        train_count = max(1, int(total * train_ratio))
        val_count = max(1, int(total * val_ratio))

        if train_count + val_count > total:
            val_count = max(0, total - train_count)

        train_selected = set(random.sample(list(images), min(train_count, len(images))))
        images -= train_selected
        val_selected = set(random.sample(list(images), min(val_count, len(images))))
        test_selected = images - val_selected

        train_files |= train_selected
        val_files |= val_selected
        test_files |= test_selected

        # Имя класса из predefined
        class_name = next((n for n, info in predefined_classes.items()
                          if info['id_1based'] == class_label), f"class_{class_label}")

        print(f"  [{class_label:2d}] {class_name[:30]:<30s}: "
              f"Train={len(train_selected):3d}, Val={len(val_selected):3d}, Test={len(test_selected):3d}")

        for img in train_selected | val_selected | test_selected:
            for cls_set in class_to_images.values():
                cls_set.discard(img)

    print(f"\nИТОГО: Train={len(train_files)}, Val={len(val_files)}, Test={len(test_files)}")

    # Подготовка labels с перенумерацией
    labels_dir_dst = output_dir / 'labels'
    
    # Временная папка
    raw_labels = output_dir / '_raw_labels'
    if raw_labels.exists():
        shutil.rmtree(raw_labels)
    raw_labels.mkdir(parents=True, exist_ok=True)
    
    print(f"\nКопирование лейблов...")
    for img in train_files | val_files | test_files:
        label_file = labels_dir / (Path(img).stem + '.txt')
        if label_file.exists():
            dst = raw_labels / (Path(img).stem + '.txt')
            if not dst.exists():
                shutil.copy2(label_file, dst)

    print(f"  Копирование с перенумерацией в YOLO 0-based формат:")
    copied, skipped = remap_labels(raw_labels, labels_dir_dst, id_to_yolo_idx)
    print(f"  ✓ Обработано {copied} файлов, пропущено строк: {skipped}")

    shutil.rmtree(raw_labels)

    # Запись txt-файлов с абсолютными путями
    def write_split(file_path: Path, file_set: set):
        with open(file_path, 'w', encoding='utf-8') as f:
            for img in sorted(file_set):
                abs_path = str(images_dir / img).replace('\\', '/')
                f.write(abs_path + '\n')

    train_txt = output_dir / f'train_{task}.txt'
    val_txt = output_dir / f'val_{task}.txt'
    test_txt = output_dir / f'test_{task}.txt'

    write_split(train_txt, train_files)
    write_split(val_txt, val_files)
    write_split(test_txt, test_files)

    # data.yaml — ВСЕГДА с полной схемой из predefined_classes.txt
    nc = len(predefined_classes)
    class_names = list(predefined_classes.keys())

    data_yaml = {
        'path': str(output_dir).replace('\\', '/'),
        'train': str(train_txt).replace('\\', '/'),
        'val': str(val_txt).replace('\\', '/'),
        'test': str(test_txt).replace('\\', '/'),
        'nc': nc,
        'names': class_names,
    }

    yaml_name = f'data_{task}.yaml'
    yaml_path = output_dir / yaml_name

    with open(yaml_path, 'w', encoding='utf-8') as f:
        yaml.safe_dump(data_yaml, f, sort_keys=False, allow_unicode=True)

    print(f"\n✓ Создан {yaml_name}: {yaml_path}")
    print(f"✓ nc={nc}, классов: {nc}")
    print(f"✓ Классы (0..{nc-1}): {class_names}")

    return yaml_path


def main():
    parser = argparse.ArgumentParser(
        description='Создание train/val/test сплита с учётом predefined_classes.txt'
    )

    parser.add_argument('--task', type=str, required=True,
                        choices=['detect', 'segment', 'both'])
    parser.add_argument('--images-dir', type=str, required=True)
    parser.add_argument('--labels-dir', type=str, required=True)
    parser.add_argument('--polygons-dir', type=str, default=None)
    parser.add_argument('--output-dir', type=str, required=True)
    parser.add_argument('--predefined-classes', type=str, default='predefined_classes.txt',
                        help='Путь к файлу с полным списком классов (по умолчанию: predefined_classes.txt)')
    parser.add_argument('--train-ratio', type=float, default=0.8)
    parser.add_argument('--val-ratio', type=float, default=0.1)
    parser.add_argument('--seed', type=int, default=42)

    args = parser.parse_args()

    # Загружаем полную схему классов
    predefined_path = Path(args.predefined_classes).resolve()
    if not predefined_path.exists():
        raise FileNotFoundError(f"Файл классов не найден: {predefined_path}\n"
                                f"Создайте его (по одному классу на строку) или укажите путь через --predefined-classes")

    predefined_classes = load_predefined_classes(predefined_path)
    print(f"📖 Загружена схема классов: {predefined_path}")
    print(f"   Всего классов: {len(predefined_classes)}")

    if args.task == 'both':
        if not args.polygons_dir:
            raise ValueError("Для task=both необходимо указать --polygons-dir")

        create_split(
            images_dir=Path(args.images_dir),
            labels_dir=Path(args.labels_dir),
            output_dir=Path(args.output_dir),
            task='detect',
            predefined_classes=predefined_classes,
            train_ratio=args.train_ratio,
            val_ratio=args.val_ratio,
            seed=args.seed,
        )

        create_split(
            images_dir=Path(args.images_dir),
            labels_dir=Path(args.polygons_dir),
            output_dir=Path(args.output_dir),
            task='segment',
            predefined_classes=predefined_classes,
            train_ratio=args.train_ratio,
            val_ratio=args.val_ratio,
            seed=args.seed,
        )

        print(f"\n{'='*60}")
        print("✅ ОБА SPLITS СОЗДАНЫ!")
        print(f"{'='*60}")

    else:
        create_split(
            images_dir=Path(args.images_dir),
            labels_dir=Path(args.labels_dir),
            output_dir=Path(args.output_dir),
            task=args.task,
            predefined_classes=predefined_classes,
            train_ratio=args.train_ratio,
            val_ratio=args.val_ratio,
            seed=args.seed,
        )

        print(f"\n{'='*60}")
        print("✅ SPLT СОЗДАН!")
        print(f"{'='*60}")
        yaml_name = f'data_{args.task}.yaml'
        if args.task == 'detect':
            print(f"Запуск: python scripts/train_yolo.py --data {Path(args.output_dir) / yaml_name} --workers 0")
        else:
            print(f"Запуск: python scripts/train_yolo_seg.py --data {Path(args.output_dir) / yaml_name} --workers 0")


if __name__ == '__main__':
    main()