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
import random
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import pandas as pd
import torch.nn.functional as F
import plotly.express as px

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)

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")
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
TARGET_K = 10  # Enforce your target 10 metabolic domains

class MSI_Structural_Encoder(nn.Module):
    def __init__(self):
        super(MSI_Structural_Encoder, self).__init__()
        # Matches the high-fidelity GroupNorm architecture
        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):
        h = self.features(x).squeeze(-1).squeeze(-1)
        return h

def load_trained_encoder(path):
    if not os.path.exists(path):
        raise FileNotFoundError(f"Missing model checkpoint: {path}")
    
    model = MSI_Structural_Encoder()
    # Loading only the feature extractor weights
    model.features.load_state_dict(torch.load(path, map_location=DEVICE))
    model.to(DEVICE)
    model.eval()
    return model

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((256, 256)),
            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("L")
        return self.transform(image), self.img_names[idx]

if __name__ == "__main__":
    print("Loading specialized custom MSI Structural Encoder...")
    encoder = load_trained_encoder(ENCODER_PATH)
    dataset = MSIDatasetInference(IMG_DIR)
    loader = DataLoader(dataset, batch_size=1, shuffle=False)

    features = []
    ion_names = []

    print("Extracting and L2-normalizing features from custom manifold...")
    with torch.no_grad():
        for img, name in loader:
            img = img.to(DEVICE)
            feat = encoder(img)
            # Re-enforce unit hypersphere mapping
            feat = F.normalize(feat, p=2, dim=1)
            features.append(feat.cpu().numpy().flatten())
            ion_names.append(name[0])

    features = np.array(features)

    # Dimensionality Reduction Optimization
    pca_km = PCA(n_components=min(50, len(features)), random_state=42)
    features_pca_all = pca_km.fit_transform(features)

    # FIX: Slice out the first principal component (index 0) 
    # if it continues to hold the global brightness/intensity bias
    features_pca_km = features_pca_all[:, 1:] 
    
    # CRITICAL FIX: Re-normalize strictly for K-Means to ensure objective Cosine geometry
    features_pca_km_norm = features_pca_km / np.linalg.norm(features_pca_km, axis=1, keepdims=True)

    # ENFORCE EXACTLY 10 METABOLIC DOMAINS
    print(f"Executing K-Means clustering strictly constrained to K={TARGET_K}...")
    clusterer = KMeans(n_clusters=TARGET_K, init='k-means++', n_init=20, random_state=42)
    labels = clusterer.fit_predict(features_pca_km_norm)

    print("Projecting t-SNE map...")
    tsne = TSNE(
        n_components=2, 
        perplexity=15,             
        early_exaggeration=24.0,   # Boosted significantly to force tighter, more definitive visual clusters
        learning_rate='auto',
        init='pca',                
        random_state=42
    )
    # Feed the unnormalized PCA vectors to t-SNE to allow intensity variance to separate the visual blobs
    embeddings_2d = tsne.fit_transform(features_pca_km)

    # Plotting
    plt.figure(figsize=(10, 8))
    scatter = plt.scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], c=labels, cmap='tab10', alpha=0.9, edgecolors='none', s=45)
    cbar = plt.colorbar(scatter, ticks=range(TARGET_K))
    cbar.set_label('Metabolic Domain ID (K=10)')
    
    plt.title("Objectively Segmented Leaf Manifold (Custom Encoder)")
    plt.xlabel("t-SNE Dimension 1")
    plt.ylabel("t-SNE Dimension 2")
    plt.grid(True, linestyle='--', alpha=0.3)
    
    plt.savefig(os.path.join(BASE_DIR, "..", "msi_tsne_corregido.png"), dpi=300, bbox_inches='tight')
    plt.close()

    # INTERACTIVE PLOTLY EXPORT
    print("Exporting Interactive Plotly t-SNE...")
    mz_values = [float(n.replace("ion_", "").replace(".png", "")) for n in ion_names]
    df_plot = pd.DataFrame({
        't-SNE Dimension 1': embeddings_2d[:, 0],
        't-SNE Dimension 2': embeddings_2d[:, 1],
        'Cluster': [str(lbl) for lbl in labels],
        'm/z': mz_values,
        'Filename': ion_names
    })
    
    fig = px.scatter(
        df_plot,
        x='t-SNE Dimension 1',
        y='t-SNE Dimension 2',
        color='Cluster',
        hover_data=['m/z', 'Filename'],
        title="Interactive Segmented Leaf Manifold (Custom Encoder)",
        category_orders={"Cluster": [str(i) for i in range(TARGET_K)]}
    )
    
    plotly_out = os.path.join(BASE_DIR, "..", "figures", "msi_tsne_interactive.html")
    os.makedirs(os.path.dirname(plotly_out), exist_ok=True)
    fig.write_html(plotly_out)
    print(f"Interactive t-SNE saved to: {plotly_out}")

    # Save outputs
    print("Exporting synchronized cluster assignments...")
    df_results = pd.DataFrame({'ion_image': ion_names, 'cluster': labels})
    df_results['mz'] = df_results['ion_image'].str.extract(r'(\d+\.\d+)').astype(float)
    df_results.to_csv(os.path.join(BASE_DIR, "..", "msi_clusters_results.csv"), index=False)
    
    print("Process complete.")
