import os
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import numpy as np
import random
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.metrics import silhouette_score, davies_bouldin_score, calinski_harabasz_score
import matplotlib.pyplot as plt
import torch.nn.functional as F

# 0. REPRODUCIBILITY
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)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)

# CONFIGURATION
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
IMG_DIR = os.path.join(BASE_DIR, "..", "datos_para_ai")
MODEL_PATH = os.path.join(BASE_DIR, "..", "models", "msi_encoder_trained.pth")
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
K_RANGE = range(2, 21) # Testing from 2 to 20 clusters

class MSI_Structural_Encoder(nn.Module):
    def __init__(self):
        super(MSI_Structural_Encoder, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=5, stride=2, padding=2, bias=False),
            nn.GroupNorm(4, 32),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1, bias=False),
            nn.GroupNorm(8, 64),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=False),
            nn.GroupNorm(16, 128),
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1, bias=False),
            nn.GroupNorm(32, 256),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1))
        )

    def forward(self, x):
        return self.features(x).squeeze(-1).squeeze(-1)

# 1. LOAD MODEL AND EXTRACT FEATURES
def extract_features():
    model = MSI_Structural_Encoder()
    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(f"No se encontró el modelo en {MODEL_PATH}")
    model.features.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
    model.to(DEVICE)
    model.eval()

    transform = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
    ])

    features = []
    if not os.path.exists(IMG_DIR):
        raise FileNotFoundError(f"No se encontró el directorio de imágenes en {IMG_DIR}")
    img_names = sorted([f for f in os.listdir(IMG_DIR) if f.endswith('.png')])
    
    print(f"Extracting features from {len(img_names)} images on {DEVICE}...")
    with torch.no_grad():
        for i, name in enumerate(img_names):
            if i % 100 == 0: print(f"Processing image {i}/{len(img_names)}...")
            img_path = os.path.join(IMG_DIR, name)
            img = Image.open(img_path).convert("L")
            img_t = transform(img).unsqueeze(0).to(DEVICE)
            feat = model(img_t)
            # CRITICAL FIX: Project explicitly onto the unit hypersphere
            feat_norm = F.normalize(feat, p=2, dim=1)
            features.append(feat_norm.cpu().numpy().flatten())
    
    return np.array(features)

# 2. EVALUATE K-MEANS
if __name__ == "__main__":
    X = extract_features()
    
    pca = PCA(n_components=min(50, X.shape[1]), random_state=42) 
    X_pca = pca.fit_transform(X)
    # CRITICAL FIX: Re-normalize after PCA to restore unit hypersphere geometry for Spherical K-Means equivalent
    X_pca = X_pca / np.linalg.norm(X_pca, axis=1, keepdims=True)

    inertias = []
    silhouettes = []
    davies_bouldin = []
    calinski_harabasz = []

    print("Evaluating different K values...")
    for k in K_RANGE:
        print(f"Calculating for K={k}...")
        kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
        labels = kmeans.fit_predict(X_pca)
        
        inertias.append(kmeans.inertia_)
        # CRITICAL FIX: Evaluate using cosine distance on the un-reduced features
        silhouettes.append(silhouette_score(X, labels, metric='cosine'))
        davies_bouldin.append(davies_bouldin_score(X_pca, labels))
        calinski_harabasz.append(calinski_harabasz_score(X_pca, labels))

    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))

    ax1.plot(K_RANGE, inertias, 'o-', color='tab:red')
    ax1.set_title('Elbow Method (Inertia)')
    ax1.set_xlabel('Number of Clusters (k)')
    ax1.set_ylabel('Sum of Squared Errors')
    ax1.grid(True)

    ax2.plot(K_RANGE, silhouettes, 's-', color='tab:blue')
    ax2.set_title('Silhouette Score (Higher is better)')
    ax2.set_xlabel('Number of Clusters (k)')
    ax2.set_ylabel('Score')
    ax2.grid(True)

    ax3.plot(K_RANGE, davies_bouldin, '^-', color='tab:green')
    ax3.set_title('Davies-Bouldin Index (Lower is better)')
    ax3.set_xlabel('Number of Clusters (k)')
    ax3.set_ylabel('Score')
    ax3.grid(True)

    ax4.plot(K_RANGE, calinski_harabasz, 'd-', color='tab:orange')
    ax4.set_title('Calinski-Harabasz Index (Higher is better)')
    ax4.set_xlabel('Number of Clusters (k)')
    ax4.set_ylabel('Score')
    ax4.grid(True)

    plt.tight_layout()
    output_plot = os.path.join(BASE_DIR, "..", "figures", "cluster_validation_detailed_plot.png")
    os.makedirs(os.path.dirname(output_plot), exist_ok=True)
    plt.savefig(output_plot)
    print(f"Plot saved as '{output_plot}'")

    idx_10 = list(K_RANGE).index(10)
    print(f"\n--- RESULTS FOR K=10 ---")
    print(f"Silhouette Score: {silhouettes[idx_10]:.4f}")
    print(f"Davies-Bouldin Index: {davies_bouldin[idx_10]:.4f}")
    print(f"Calinski-Harabasz Index: {calinski_harabasz[idx_10]:.4f}")
    
    print(f"\n--- MATHEMATICAL SUGGESTIONS ---")
    print(f"Best K (Silhouette): {K_RANGE[np.argmax(silhouettes)]}")
    print(f"Best K (Davies-Bouldin): {K_RANGE[np.argmin(davies_bouldin)]}")
    print(f"Best K (Calinski-Harabasz): {K_RANGE[np.argmax(calinski_harabasz)]}")