import os
import random
import warnings
from pathlib import Path
from PIL import Image
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
import timm

# Подавление предупреждений цветовых профилей
warnings.filterwarnings("ignore", message=".*Unknown Adobe color transform code.*")

# Настройки окружения
DATA_ROOT = Path("/home/zin/projects/Thesis/NFS/Thesis/Emoset/EmoSet-118K")
BATCH_SIZE = 64
EPOCHS = 30
LR = 5e-5
NUM_WORKERS = 62
PATIENCE = 7

# Маппинг классов
CLASS_MAPPING = {
    "amusement": 0, "anger": 1, "awe": 2, "contentment": 3,
    "disgust": 4, "excitement": 5, "fear": 6, "sadness": 7
}

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

# Фиксация генераторов псевдослучайных чисел
def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)

set_seed()

# Инициализация структур данных
class EmoSetDataset(Dataset):
    def __init__(self, root: Path | str, split: str, transform=None):
        self.root = Path(root) / split
        self.df = pd.read_csv(self.root / "labels.csv")
        self.transform = transform
        
        # Фильтрация датафрейма
        self.df = self.df[self.df["label"].isin(CLASS_MAPPING.keys())].reset_index(drop=True)

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

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        img_path = self.root / "images" / row["filename"]

        try:
            img = Image.open(img_path).convert("RGB")
        except Exception:
            img = Image.new("RGB", (256, 256), (0, 0, 0))

        if self.transform:
            img_tensor = self.transform(img)
        else:
            img_tensor = T.ToTensor()(img)
            
        label_idx = CLASS_MAPPING[row["label"]]
        return img_tensor, label_idx

# Трансформации
base_tf = [
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
]

train_transform = T.Compose([
    T.Resize(256, antialias=True),
    T.RandomCrop(224),
    T.RandomHorizontalFlip(),
    *base_tf
])

val_transform = T.Compose([
    T.Resize(256, antialias=True),
    T.CenterCrop(224),
    *base_tf
])

train_ds = EmoSetDataset(DATA_ROOT, "train", transform=train_transform)
val_ds   = EmoSetDataset(DATA_ROOT, "val", transform=val_transform)

train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=True)
val_loader   = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)

# Инициализация модели и оптимизатора
model = timm.create_model("resnet50", pretrained=True, num_classes=8, drop_rate=0.3)
model.to(DEVICE)

criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-3)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)

# Логика эпохи обучения
def train_epoch(current_model, loader):
    current_model.train()
    total_loss, correct_preds, total_samples = 0.0, 0, 0

    for imgs, labels in tqdm(loader, desc="Тренировка", leave=False, smoothing=0):
        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)

        optimizer.zero_grad(set_to_none=True)
        logits = current_model(imgs)
        loss = criterion(logits, labels)

        loss.backward()
        optimizer.step()

        total_loss += loss.item() * imgs.size(0)
        preds = logits.argmax(dim=1)
        correct_preds += (preds == labels).sum().item()
        total_samples += labels.size(0)

    return total_loss / total_samples, correct_preds / total_samples

# Логика эпохи валидации
@torch.no_grad()
def val_epoch(current_model, loader):
    current_model.eval()
    total_loss, correct_preds, total_samples = 0.0, 0, 0

    for imgs, labels in tqdm(loader, desc="Валидация", leave=False, smoothing=0):
        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)

        logits = current_model(imgs)
        loss = criterion(logits, labels)

        total_loss += loss.item() * imgs.size(0)
        preds = logits.argmax(dim=1)
        correct_preds += (preds == labels).sum().item()
        total_samples += labels.size(0)

    return total_loss / total_samples, correct_preds / total_samples

if __name__ == "__main__":
    best_val_acc = 0.0
    best_val_loss = float('inf')
    epochs_no_improve = 0
    checkpoint_path = "./emosetV2_resnet50_best.pth"

    print("Старт обучения.")

    for epoch in range(1, EPOCHS + 1):
        train_loss, train_acc = train_epoch(model, train_loader)
        val_loss, val_acc = val_epoch(model, val_loader)

        scheduler.step()

        print(f"[{epoch}/{EPOCHS}] Train Loss: {train_loss:.4f}, Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f}, Acc: {val_acc:.4f}")

        # Сохранение лучших весов по Accuracy
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), checkpoint_path)
            print(f"Сохранен чекпоинт (Acc: {best_val_acc:.4f})")

        # Оценка переобучения по Loss (Early Stopping)
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            epochs_no_improve = 0
        else:
            epochs_no_improve += 1
            if epochs_no_improve >= PATIENCE:
                print(f"Ранняя остановка: метрика валидации не улучшается {PATIENCE} эпох.")
                break

    print("Процесс завершен.")