import os
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import numpy as np
import random
import pandas as pd
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.metrics import silhouette_score, davies_bouldin_score, calinski_harabasz_score, adjusted_rand_score
import matplotlib.pyplot as plt
from scipy.stats import ttest_ind
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 = 10

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. DATA LOADING
def load_images():
    img_names = sorted([f for f in os.listdir(IMG_DIR) if f.endswith('.png')])
    raw_pixels = []
    mzs = []
    
    print(f"Loading {len(img_names)} images for validation...")
    for name in img_names:
        img_path = os.path.join(IMG_DIR, name)
        img = Image.open(img_path).convert("L")
        raw_pixels.append(np.array(img).flatten())
        
        try:
            mz = float(name.split('_')[1].replace('.png', ''))
            mzs.append(mz)
        except:
            mzs.append(0.0)
            
    return np.array(raw_pixels), np.array(mzs), img_names

def extract_simclr_features(img_names):
    model = MSI_Structural_Encoder()
    if os.path.exists(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 = []
    print(f"Extracting SimCLR features on {DEVICE}...")
    with torch.no_grad():
        for name in 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. VALIDATION TASKS
def baseline_comparison(X_raw, X_simclr):
    print("\n--- TASK 1: BASELINE COMPARISON (PCA vs SimCLR) ---")
    pca_raw = PCA(n_components=50, random_state=42)
    X_pca_raw = pca_raw.fit_transform(X_raw)
    
    pca_simclr = PCA(n_components=50, random_state=42)
    X_pca_simclr = pca_simclr.fit_transform(X_simclr)
    
    km_raw = KMeans(n_clusters=K, random_state=42, n_init=10)
    labels_raw = km_raw.fit_predict(X_pca_raw)
    
    km_sim = KMeans(n_clusters=K, random_state=42, n_init=10)
    labels_sim = km_sim.fit_predict(X_pca_simclr)
    
    metrics = {
        "Method": ["Raw PCA", "SimCLR Encoder"],
        "Silhouette": [silhouette_score(X_pca_raw, labels_raw), silhouette_score(X_pca_simclr, labels_sim)],
        "Davies-Bouldin": [davies_bouldin_score(X_pca_raw, labels_raw), davies_bouldin_score(X_pca_simclr, labels_sim)],
        "Calinski-Harabasz": [calinski_harabasz_score(X_pca_raw, labels_raw), calinski_harabasz_score(X_pca_simclr, labels_sim)]
    }
    
    df_metrics = pd.DataFrame(metrics)
    print(df_metrics.to_string(index=False))
    return df_metrics

def stability_analysis(X_features):
    print("\n--- TASK 2: STABILITY ANALYSIS (Multiple Seeds) ---")
    seeds = [10, 20, 30, 42, 50, 60, 70, 80, 90, 100]
    all_labels = []
    metrics_list = []
    
    pca = PCA(n_components=50, random_state=42)
    X_pca = pca.fit_transform(X_features)

    for seed in seeds:
        km = KMeans(n_clusters=K, random_state=seed, n_init=10)
        labels = km.fit_predict(X_pca)
        all_labels.append(labels)
        metrics_list.append(silhouette_score(X_pca, labels))
    
    ari_scores = []
    for i in range(len(all_labels)):
        for j in range(i + 1, len(all_labels)):
            ari_scores.append(adjusted_rand_score(all_labels[i], all_labels[j]))
            
    print(f"Mean Silhouette Score: {np.mean(metrics_list):.4f} ± {np.std(metrics_list):.4f}")
    print(f"Mean Adjusted Rand Index (ARI) between runs: {np.mean(ari_scores):.4f} (1.0 is perfect stability)")

def isotope_validation(X_features, mzs):
    print("\n--- TASK 3: ISOTOPE CO-LOCALIZATION VALIDATION ---")
    X_norm = X_features / np.linalg.norm(X_features, axis=1, keepdims=True)
    
    isotope_dist = []
    random_dist = []
    
    for i in range(len(mzs)):
        for j in range(i + 1, len(mzs)):
            diff = abs(mzs[i] - mzs[j])
            if 0.95 <= diff <= 1.05:
                dist = np.linalg.norm(X_norm[i] - X_norm[j])
                isotope_dist.append(dist)
    
    if len(isotope_dist) > 0:
        for _ in range(len(isotope_dist)):
            idx1 = random.randint(0, len(mzs)-1)
            idx2 = random.randint(0, len(mzs)-1)
            if idx1 != idx2:
                dist = np.linalg.norm(X_norm[idx1] - X_norm[idx2])
                random_dist.append(dist)
                
        print(f"Number of isotopic pairs found: {len(isotope_dist)}")
        print(f"Mean distance (Isotopes): {np.mean(isotope_dist):.4f}")
        print(f"Mean distance (Random): {np.mean(random_dist):.4f}")
        
        t_stat, p_val = ttest_ind(isotope_dist, random_dist)
        print(f"T-test p-value: {p_val:.4e} (Significant if < 0.05)")
        
        plt.figure(figsize=(8, 6))
        plt.boxplot([isotope_dist, random_dist], tick_labels=['Isotopic Pairs', 'Random Pairs'])
        plt.title('Latent Space Distances: Isotopes vs Random')
        plt.ylabel('Euclidean Distance (Normalized Latent Space)')
        plt.savefig(os.path.join(BASE_DIR, "..", "figures", "isotope_validation_plot.png"))
        print(f"Isotope validation plot saved to figures/isotope_validation_plot.png")

if __name__ == "__main__":
    if not os.path.exists(IMG_DIR):
        print(f"Error: IMG_DIR {IMG_DIR} not found. Generate ion images first.")
    else:
        X_raw, mzs, img_names = load_images()
        X_simclr = extract_simclr_features(img_names)
        
        baseline_comparison(X_raw, X_simclr)
        stability_analysis(X_simclr)
        isotope_validation(X_simclr, mzs)