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.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from sklearn.manifold import TSNE
import pandas as pd
import matplotlib.pyplot as plt

# 1. CONFIGURACIÓN
IMG_DIR = "datos_para_ai"
OUTPUT_DIR = "pruebas_clusters"
N_CLUSTERS = 10 
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

os.makedirs(OUTPUT_DIR, exist_ok=True)

# 2. CARGAR MODELO PRE-ENTRENADO (Directo de PyTorch)
def get_pretrained_extractor():
    # Usamos los pesos por defecto de ImageNet
    model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.DEFAULT)
    model.classifier = nn.Identity() # Quitamos la capa de clasificación
    model.to(DEVICE)
    model.eval()
    return model

# 3. DATASET CON NORMALIZACIÓN ESTÁNDAR
class MSIDatasetDirect(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(),
            # Normalización obligatoria para modelos pre-entrenados
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    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]

if __name__ == "__main__":
    print(f"Extrayendo características con red pre-entrenada en {DEVICE}...")
    model = get_pretrained_extractor()
    dataset = MSIDatasetDirect(IMG_DIR)
    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 = model(img)
            # Aplanar el vector de características (1280 dimensiones para EffNet-B0)
            features.append(feat.cpu().numpy().flatten())
            ion_names.append(name[0])

    features = np.array(features)
    print(f"Características extraídas: {features.shape}")

    # 4. REDUCCIÓN DE DIMENSIONALIDAD (PCA)
    # Esto ayuda a que el clustering no se pierda en dimensiones irrelevantes
    print("Reduciendo ruido con PCA...")
    features_scaled = StandardScaler().fit_transform(features)
    pca = PCA(n_components=0.95) # Mantenemos el 95% de la varianza
    features_pca = pca.fit_transform(features_scaled)
    print(f"Dimensiones tras PCA: {features_pca.shape[1]}")

    # 5. CLUSTERING K-MEANS
    print(f"Agrupando en {N_CLUSTERS} clusters...")
    km = KMeans(n_clusters=N_CLUSTERS, n_init=20, random_state=42)
    labels = km.fit_predict(features_pca)

    # 6. VISUALIZACIÓN (t-SNE)
    print("Generando visualización t-SNE para verificar consistencia...")
    tsne = TSNE(n_components=2, random_state=42)
    features_tsne = tsne.fit_transform(features_pca)
    
    plt.figure(figsize=(10, 8))
    scatter = plt.scatter(features_tsne[:, 0], features_tsne[:, 1], c=labels, cmap='tab10', alpha=0.7)
    plt.colorbar(scatter, ticks=range(N_CLUSTERS), label='Cluster ID')
    plt.title(f"Visualización de Clusters (K={N_CLUSTERS}) - Directo")
    plt.xlabel("t-SNE 1")
    plt.ylabel("t-SNE 2")
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.savefig(os.path.join(OUTPUT_DIR, "msi_clusters_direct_plot.png"))
    print(f"Gráfico de consistencia guardado en '{OUTPUT_DIR}/msi_clusters_direct_plot.png'.")

    # 7. GUARDAR Y REPORTAR
    df = pd.DataFrame({'ion_image': ion_names, 'cluster': labels})
    df['mz'] = df['ion_image'].str.extract(r'(\d+\.\d+)').astype(float)
    df.to_csv(os.path.join(OUTPUT_DIR, "msi_clusters_direct_results.csv"), index=False)
    
    # Crear un resumen de cuántos iones hay por cluster
    print("\nResumen de Clusters (Consistente con el CSV y el Gráfico):")
    counts = df['cluster'].value_counts().sort_index()
    print(counts)
    
    # Guardar también el resumen para evitar confusiones con otros archivos
    summary_df = counts.reset_index()
    summary_df.columns = ['Cluster_ID', 'Ion_Count']
    summary_df.to_csv(os.path.join(OUTPUT_DIR, "msi_clusters_direct_summary.csv"), index=False)
    
    print(f"\nResultados guardados en '{OUTPUT_DIR}/msi_clusters_direct_results.csv' y '{OUTPUT_DIR}/msi_clusters_direct_summary.csv'.")
    print(f"Usa '{OUTPUT_DIR}/msi_clusters_direct_plot.png' para validar visualmente los puntos.")