import os
from pathlib import Path
from PIL import Image
import pandas as pd
import numpy as np
from tqdm import tqdm

import torch
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T
import timm

import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.manifold import TSNE

# Настройки путей для медиа
MEDIA_DIR = Path("scripts/media")
MEDIA_DIR.mkdir(parents=True, exist_ok=True)

# Конфигурация путей для инференса и кэширования векторов
DATA_ROOT = Path("./NFS/Thesis/Emoset/EmoSet-118K")
MODEL_PATH = Path("./src/emoset_resnet50_best.pth")

BATCH_SIZE = 128
NUM_WORKERS = 32

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Вычисления перенесены на: {device}")

class EmoSetFeatureDataset(Dataset):
    def __init__(self, root: Path | str, split: str):
        self.root = Path(root) / split
        self.df = pd.read_csv(self.root / "labels.csv")

        self.labels = sorted(self.df["label"].unique())
        self.label2idx = {l: i for i, l in enumerate(self.labels)}
        self.idx2label = {i: l for l, i in self.label2idx.items()}

        # Для экстракции признаков аугментация отключена, используется строгий CenterCrop
        self.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])
        ])

    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", (224, 224), (0, 0, 0))
            
        img_tensor = self.transform(img)
        label_idx = self.label2idx[row["label"]]
        
        return img_tensor, label_idx

def plot_tsne(embeddings, labels, idx2label, sample_limit=3000):
    """Генерация t-SNE графика для диссертации"""
    print(f"Построение t-SNE проекции для {sample_limit} сэмплов...")
    
    tsne_model = TSNE(n_components=2, perplexity=30, random_state=42)
    embeddings_2d = tsne_model.fit_transform(embeddings[:sample_limit])
    labels_subset = labels[:sample_limit]

    plt.figure(figsize=(12, 9))
    
    # Используем более академическую палитру
    scatter = plt.scatter(
        embeddings_2d[:, 0],
        embeddings_2d[:, 1],
        c=labels_subset,
        cmap="Set2", # Set2 лучше различается при печати
        alpha=0.7,
        s=20,
        edgecolors='w', 
        linewidths=0.5
    )

    # Формирование легенды
    handles, _ = scatter.legend_elements()
    legend_labels = [idx2label[i] for i in range(len(idx2label))]
    
    # Размещение легенды снаружи графика, чтобы не перекрывать данные
    plt.legend(handles, legend_labels, title="Эмоциональные классы", 
               bbox_to_anchor=(1.05, 1), loc='upper left')

    plt.title("2D проекция скрытого пространства признаков (t-SNE)", pad=20, fontsize=14)
    plt.xlabel("Первая главная компонента (t-SNE 1)", fontsize=12)
    plt.ylabel("Вторая главная компонента (t-SNE 2)", fontsize=12)
    plt.grid(True, linestyle='--', alpha=0.3)
    
    plt.tight_layout()
    plot_path = MEDIA_DIR / "tsne_embeddings.png"
    plt.savefig(plot_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"[INFO] График t-SNE сохранен в: {plot_path}")


if __name__ == "__main__":
    test_ds = EmoSetFeatureDataset(DATA_ROOT, "test")
    test_loader = DataLoader(
        test_ds,
        batch_size=BATCH_SIZE,
        shuffle=False,  # Отключение шаффла для строгого соответствия индексов
        num_workers=NUM_WORKERS,
        pin_memory=True
    )

    print(f"Подготовлено для извлечения: {len(test_ds)} файлов.")
    
    # Инициализация модели и загрузка лучших весов
    feature_extractor = timm.create_model(
        "resnet50",
        pretrained=False,
        num_classes=len(test_ds.labels)
    )

    try:
        checkpoint = torch.load(MODEL_PATH, map_location=device)
        feature_extractor.load_state_dict(checkpoint)
        print("Веса модели успешно загружены.")
    except Exception as e:
        print(f"Ошибка загрузки весов: {e}. Убедитесь, что модель обучена.")
        exit(1)

    # Удаление классификационного слоя (fc)
    feature_extractor.reset_classifier(0)
    feature_extractor.to(device)
    feature_extractor.eval()
    
    print("Слой классификации удален. Модель готова к экстракции.")
    
    extracted_embeddings = []
    extracted_labels = []

    print("Старт пакетной экстракции признаков...")

    with torch.no_grad():
        for imgs, labels in tqdm(test_loader, desc="Экстракция"):
            imgs = imgs.to(device)

            # Получение вектора [BATCH_SIZE, 2048]
            embeddings_batch = feature_extractor(imgs)
            
            extracted_embeddings.append(embeddings_batch.cpu().numpy())
            extracted_labels.append(labels.numpy())

    # Агрегация батчей в единые массивы
    np_embeddings = np.concatenate(extracted_embeddings, axis=0)
    np_labels = np.concatenate(extracted_labels, axis=0)

    print(f"Размерность матрицы признаков: {np_embeddings.shape}")

    # Сохранение артефактов
    np.save("./src/emoset_test_embeddings.npy", np_embeddings)
    np.save("./src/emoset_test_labels.npy", np_labels)
    print("Матрицы успешно экспортированы в .npy файлы.")
    
    # Генерация медиа для диссертации
    plot_tsne(np_embeddings, np_labels, test_ds.idx2label, sample_limit=3000)
    
    print("Процесс полностью завершен.")