import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torchvision import models, transforms
from PIL import Image
import numpy as np
from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import pandas as pd

# 1. CONFIGURACIÓN
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
IMG_DIR = os.path.join(BASE_DIR, "../..", "datos_para_ai")
ENCODER_PATH = os.path.join(BASE_DIR, "../..", "models", "msi_encoder_trained.pth")
N_CLUSTERS = 10 
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 2. CARGAR EL MODELO ENTRENADO
def load_trained_encoder(path):
    if not os.path.exists(path):
        raise FileNotFoundError(f"No se encontró el modelo en {path}")
    
    model = models.efficientnet_b0()
    model.classifier = nn.Identity()
    
    state_dict = torch.load(path, map_location=DEVICE)
    model.load_state_dict(state_dict)
    model.to(DEVICE)
    model.eval()
    return model

# 3. DATASET PARA INFERENCIA
class MSIDatasetInference(Dataset):
    def __init__(self, img_dir):
        self.img_dir = img_dir
        self.img_names = sorted([f for f in os.listdir(img_dir) if f.endswith('.png')])
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
        ])

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

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_names[idx])
        image = Image.open(img_path).convert("RGB")
        return self.transform(image), self.img_names[idx]

# --- PROCESO ---
if __name__ == "__main__":
    print("Extrayendo características con el encoder...")
    try:
        encoder = load_trained_encoder(ENCODER_PATH)
    except Exception as e:
        print(f"Error al cargar el modelo: {e}")
        exit(1)

    dataset = MSIDatasetInference(IMG_DIR)
    if len(dataset) == 0:
        print(f"No se encontraron imágenes en {IMG_DIR}")
        exit(1)
        
    loader = DataLoader(dataset, batch_size=1, shuffle=False)

    features = []
    ion_names = []

    with torch.no_grad():
        for img, name in loader:
            img = img.to(DEVICE)
            feat = encoder(img)
            # Normalización L2
            feat = feat / feat.norm(p=2, dim=1, keepdim=True)
            features.append(feat.cpu().numpy().flatten())
            ion_names.append(name[0])

    features = np.array(features)

    # 4. CLUSTERING K-MEANS
    print(f"Realizando KMeans con {N_CLUSTERS} clusters...")
    km = KMeans(n_clusters=N_CLUSTERS, n_init=20, random_state=42)
    labels = km.fit_predict(features)

    # 5. VISUALIZACIÓN t-SNE CORREGIDA
    print("Generando visualización t-SNE con escala corregida...")
    tsne = TSNE(n_components=2, perplexity=30, random_state=42)
    embeddings_2d = tsne.fit_transform(features)

    plt.figure(figsize=(10, 8))
    # Aplicando tu corrección específica:
    scatter = plt.scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], c=labels, cmap='tab10', alpha=0.8, edgecolors='w', s=50)
    
    # Configuración de la barra de colores con los ticks exactos para 10 clusters
    cbar = plt.colorbar(scatter, ticks=range(N_CLUSTERS))
    cbar.set_label('Cluster ID')
    
    plt.title(f"t-SNE Visualization of MSI Ions (K={N_CLUSTERS})")
    plt.xlabel("t-SNE dimension 1")
    plt.ylabel("t-SNE dimension 2")
    plt.grid(True, linestyle='--', alpha=0.5)
    
    output_file = "msi_tsne_corregido.png"
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    # plt.show()
    
    print(f"Visualización guardada exitosamente en '{output_file}'.")