import os
import random
import warnings
from collections import defaultdict
from pathlib import Path
from PIL import Image, ImageFile
import pandas as pd
import numpy as np
from tqdm import tqdm

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T
from torch.amp import autocast, GradScaler
import timm

# Подавление предупреждений и защита от битых "хвостов" JPEG
warnings.filterwarnings("ignore")
ImageFile.LOAD_TRUNCATED_IMAGES = True

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Устройство: {DEVICE}")

# --- ПУТИ ---
TRAIN_ROOT = Path("./dataset/Original-2.41M")
ANCHOR_118K_ROOT = Path("./NFS/Thesis/Emoset/EmoSet-118K/train") # ЯКОРЬ (Чистые данные для обучения)
VAL_118K_ROOT = Path("./NFS/Thesis/Emoset/EmoSet-118K/val")

SAVE_MODEL_PATH = Path("./src/emosetV2_resnet50_finetuned_2_41M.pth")
RESUME_CHECKPOINT = Path("./src/finetuneV2_resume.pth")
PRETRAINED_PATH = Path("./src/emosetV2_resnet50_best.pth")

CLASS_MAPPING = {
    "amusement": 0, "anger": 1, "awe": 2, "contentment": 3,
    "disgust": 4, "excitement": 5, "fear": 6, "sadness": 7
}

# --- НАСТРОЙКИ ---
TOTAL_BATCH_SIZE = 64
BATCH_NOISY = 48  # 75% батча - новые данные 2.41M
BATCH_ANCHOR = 16 # 25% батча - чистые якорные данные 118K

EPOCHS_PER_FOLDER = 15
PATIENCE = 5
LR = 1e-6 
NUM_TRAIN_WORKERS = 32
NUM_VAL_WORKERS = 32

def worker_init_fn(worker_id):
    np.random.seed(np.random.get_state()[1][0] + worker_id)

# --- 1. ТРАНСФОРМАЦИИ ---
train_transform = T.Compose([
    T.Resize(256),
    T.RandomResizedCrop(224, scale=(0.8, 1.0)),
    T.RandomHorizontalFlip(),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

val_transform = T.Compose([
    T.Resize(256),
    T.CenterCrop(224),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# --- 2. ДАТАСЕТЫ ---
class ChunkTrainDataset(Dataset):
    def __init__(self, paths, transform):
        self.paths = paths
        self.transform = transform

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx):
        path = self.paths[idx]
        try:
            img = Image.open(path).convert('RGB')
            tensor = self.transform(img)
            label = CLASS_MAPPING.get(path.parts[-3].lower(), 0)
            return tensor, label
        except Exception:
            return torch.zeros((3, 224, 224)), 0

class CsvDataset(Dataset):
    def __init__(self, root, transform):
        self.root = Path(root)
        self.df = pd.read_csv(self.root / "labels.csv")
        self.transform = transform

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        path = self.root / "images" / row["filename"]
        try:
            img = Image.open(path).convert('RGB')
            tensor = self.transform(img)
            label = CLASS_MAPPING.get(row["label"].lower(), 0)
            return tensor, label
        except Exception:
            return torch.zeros((3, 224, 224)), 0

# --- 3. СБОР ДАННЫХ ---
def prepare_chunks():
    print("\nСканирование датасета 2.41M...")
    chunk_dict = defaultdict(list)
    for path in TRAIN_ROOT.rglob('*.jpg'):
        emotion = path.parts[-3].lower()
        if emotion not in CLASS_MAPPING: 
            continue
        folder_str = path.parts[-2]
        if folder_str.isdigit():
            chunk_dict[int(folder_str)].append(path)
            
    sorted_chunks = sorted(chunk_dict.keys())
    print(f"Найдено пронумерованных папок (чанков): {len(sorted_chunks)}")
    return chunk_dict, sorted_chunks
    # --- 4. ОСНОВНОЙ ЦИКЛ ОБУЧЕНИЯ ---
if __name__ == "__main__":
    chunk_dict, sorted_chunks = prepare_chunks()

    # Валидационный датасет (только чистые данные)
    val_loader = DataLoader(
        CsvDataset(VAL_118K_ROOT, val_transform), 
        batch_size=TOTAL_BATCH_SIZE, shuffle=False, 
        num_workers=NUM_VAL_WORKERS, pin_memory=True
    )

    # ЯКОРНЫЙ ЗАГРУЗЧИК (Чистые данные для подмешивания)
    # Используем prefetch_factor и persistent_workers для устранения рывков CPU
    anchor_dataset = CsvDataset(ANCHOR_118K_ROOT, train_transform)
    anchor_loader = DataLoader(
        anchor_dataset, batch_size=BATCH_ANCHOR, shuffle=True, 
        num_workers=16, pin_memory=True, drop_last=True,
        prefetch_factor=2, persistent_workers=False
    )

    # Инициализация модели
    model = timm.create_model('resnet50', pretrained=False, num_classes=8).to(DEVICE)
    if PRETRAINED_PATH.exists():
        model.load_state_dict(torch.load(PRETRAINED_PATH, map_location=DEVICE))
        print(f"Базовые веса загружены из {PRETRAINED_PATH.name}")

    # Размораживаем всю модель
    for param in model.parameters(): 
        param.requires_grad = True

    # Дифференцированный оптимизатор
    backbone_params = [p for n, p in model.named_parameters() if "fc" not in n]
    fc_params = [p for n, p in model.named_parameters() if "fc" in n]
    
    optimizer = torch.optim.AdamW([
        {'params': backbone_params, 'lr': LR},         # 1e-6: микро-шаг для основы
        {'params': fc_params, 'lr': LR * 10}           # 1e-5: шаг для классификатора
    ], weight_decay=1e-3)

    # Label Smoothing помогает игнорировать мусор в разметке 2.41M
    criterion = nn.CrossEntropyLoss(label_smoothing=0.15)
    scaler = GradScaler()

    # --- ПАРАМЕТРЫ ВОССТАНОВЛЕНИЯ ---
    start_stage = 0
    start_epoch = 1
    best_val_loss = float('inf')

    if RESUME_CHECKPOINT.exists():
        print(f"\nОбнаружен чекпоинт: {RESUME_CHECKPOINT.name}. Восстановление...")
        checkpoint = torch.load(RESUME_CHECKPOINT, map_location=DEVICE)
        model.load_state_dict(checkpoint['model_state_dict'])
        
        try: 
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        except Exception as e: 
            print(f"Оптимизатор сброшен: {e}")
            
        best_val_loss = checkpoint['best_val_loss']
        start_stage = checkpoint['stage']
        start_epoch = checkpoint['epoch'] + 1 
        print(f"Успешный запуск с ЭТАПА {start_stage + 1}, Эпохи {start_epoch}. Best Val Loss: {best_val_loss:.4f}\n")
    else:
        # --- ЗАМЕР EPOCH 0 (БАЗОВАЯ ТОЧНОСТЬ) ---
        # Выполняется только если мы начинаем с нуля
        print("\n[Проверка базовых весов перед обучением (Epoch 0)]")
        model.eval()
        val_loss, val_correct, val_total = 0.0, 0, 0
        with torch.no_grad():
            for inputs, labels in tqdm(val_loader, desc="Baseline Eval", smoothing=0):
                inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
                with autocast(device_type="cuda"):
                    outputs = model(inputs)
                    v_loss = criterion(outputs, labels)
                val_loss += v_loss.item() * inputs.size(0)
                _, pred = outputs.max(1)
                val_total += labels.size(0)
                val_correct += pred.eq(labels).sum().item()
        
        best_val_loss = val_loss / val_total
        baseline_acc = val_correct / val_total
        print(f"Стартовая точка -> Val Loss: {best_val_loss:.4f} | Val Acc: {baseline_acc:.4f}\n")

    # ВОССТАНОВЛЕНИЕ НАКОПЛЕННЫХ ДАННЫХ
    current_train_paths = []
    for s in range(start_stage):
        current_train_paths.extend(chunk_dict[sorted_chunks[s]])

    print("Старт Anchor Curriculum Learning (Смешивание чистых и шумных данных).")

    # ГЛАВНЫЙ ЦИКЛ ПО ПАПКАМ
    for stage in range(start_stage, len(sorted_chunks)):
        chunk_id = sorted_chunks[stage]
        print(f"\n{'='*50}")
        print(f"ЭТАП {stage+1}/{len(sorted_chunks)}: Добавляем папку '{chunk_id}'")
        
        # Накопление и перемешивание
        current_train_paths.extend(chunk_dict[chunk_id])
        random.shuffle(current_train_paths)
        print(f"Всего файлов (грязных) в текущем пуле: {len(current_train_paths)}")
        
        # ОСНОВНОЙ ЗАГРУЗЧИК (Грязные данные) с PREFETCH
        train_loader = DataLoader(
            ChunkTrainDataset(current_train_paths, train_transform), 
            batch_size=BATCH_NOISY, shuffle=True, 
            num_workers=NUM_TRAIN_WORKERS, pin_memory=True, 
            worker_init_fn=worker_init_fn, drop_last=True,
            prefetch_factor=4, persistent_workers=True # Устраняет рывки CPU
        )

        epochs_no_improve = 0
        first_epoch = start_epoch if stage == start_stage else 1
        
        # Инициализация итератора якорей
        anchor_iter = iter(anchor_loader)

        # ЦИКЛ ЭПОХ ДЛЯ ТЕКУЩЕГО ЭТАПА
        for epoch in range(first_epoch, EPOCHS_PER_FOLDER + 1):
            model.train()
            train_loss, train_correct, train_total = 0.0, 0, 0
            
            for noisy_inputs, noisy_labels in tqdm(train_loader, desc=f"S{stage+1}-Ep{epoch}/{EPOCHS_PER_FOLDER} [Train]", smoothing=0):
                
                # Достаем якорный чистый батч
                try:
                    anc_inputs, anc_labels = next(anchor_iter)
                except StopIteration:
                    anchor_iter = iter(anchor_loader)
                    anc_inputs, anc_labels = next(anchor_iter)
                
                # СМЕШИВАЕМ БАТЧИ (Грязные + Чистые)
                inputs = torch.cat([noisy_inputs, anc_inputs]).to(DEVICE)
                labels = torch.cat([noisy_labels, anc_labels]).to(DEVICE)
                
                optimizer.zero_grad(set_to_none=True)
                with autocast(device_type="cuda"):
                    outputs = model(inputs)
                    loss = criterion(outputs, labels)
                    
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
                
                train_loss += loss.item() * inputs.size(0)
                _, pred = outputs.max(1)
                train_total += labels.size(0)
                train_correct += pred.eq(labels).sum().item()

            # ВАЛИДАЦИЯ
            model.eval()
            val_loss, val_correct, val_total = 0.0, 0, 0
            with torch.no_grad():
                for inputs, labels in tqdm(val_loader, desc="[Val]", leave=False, smoothing=0):
                    inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
                    with autocast(device_type="cuda"):
                        outputs = model(inputs)
                        v_loss = criterion(outputs, labels)
                    val_loss += v_loss.item() * inputs.size(0)
                    _, pred = outputs.max(1)
                    val_total += labels.size(0)
                    val_correct += pred.eq(labels).sum().item()
            
            avg_train_loss = train_loss / train_total
            avg_train_acc = train_correct / train_total
            avg_val_loss = val_loss / val_total
            avg_val_acc = val_correct / val_total
            
            print(f"S{stage+1}-E{epoch} | Train L: {avg_train_loss:.4f}, Acc: {avg_train_acc:.4f} | Val L: {avg_val_loss:.4f}, Acc: {avg_val_acc:.4f}")

            # СОХРАНЕНИЕ ЛУЧШИХ ВЕСОВ
            if avg_val_loss < best_val_loss:
                best_val_loss = avg_val_loss
                epochs_no_improve = 0
                torch.save(model.state_dict(), SAVE_MODEL_PATH)
                print("--> Обновлены лучшие веса")
            else:
                epochs_no_improve += 1

            # АВАРИЙНОЕ СОХРАНЕНИЕ В КОНЦЕ ЭПОХИ
            checkpoint_state = {
                'stage': stage,
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_val_loss': best_val_loss
            }
            torch.save(checkpoint_state, RESUME_CHECKPOINT)
            os.sync() # Защита от отключения электричества
            print(f"--> Чекпоинт (Этап {stage+1}, Эпоха {epoch}) зафиксирован на диске.")

            # РАННЯЯ ОСТАНОВКА ДЛЯ ТЕКУЩЕГО ЭТАПА
            if epochs_no_improve >= PATIENCE:
                print(f"Ранняя остановка для ЭТАПА {stage+1}. Переход к следующей папке...")
                break
        
        # Сброс счетчика стартовой эпохи после прохождения восстановительного этапа
        start_epoch = 1