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
import torch.nn.functional as F

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")

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)

def load_encoder_diagnostic(path):
    model = MSI_Structural_Encoder()
    if os.path.exists(path):
        state_dict = torch.load(path, map_location=DEVICE)
        model.features.load_state_dict(state_dict)
    model.to(DEVICE)
    model.eval()
    return model

if __name__ == "__main__":
    encoder = load_encoder_diagnostic(ENCODER_PATH)
    
    # Load targets directly to test raw alignment
    img_155 = os.path.join(IMG_DIR, "ion_155.75.png")
    img_156 = os.path.join(IMG_DIR, "ion_156.07.png")
    
    transform = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor()
    ])
    
    if os.path.exists(img_155) and os.path.exists(img_156):
        # Convert to Grayscale "L"
        t_155 = transform(Image.open(img_155).convert("L")).unsqueeze(0).to(DEVICE)
        t_156 = transform(Image.open(img_156).convert("L")).unsqueeze(0).to(DEVICE)
        
        with torch.no_grad():
            f_155 = encoder(t_155)
            f_156 = encoder(t_156)
            
            # Check directional similarity
            sim = F.cosine_similarity(f_155, f_156)
            print(f"\n--- Feature Diagnostics ---")
            print(f"Cosine Similarity between Vein Ions (155.75 vs 156.07): {sim.item():.4f}")
            print("If this value is below 0.85, the loaded weight file is treating them as different shapes.")