import os
import torch
import torch.nn.functional as F
import pandas as pd
import numpy as np
from PIL import Image
from torchvision import transforms

# Copy the exact encoder structure from your scripts
import torch.nn as nn
class MSI_Structural_Encoder(nn.Module):
    def __init__(self):
        super().__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)

DEVICE = torch.device("cpu")
IMG_DIR = "datos_para_ai"
MODEL_PATH = "models/msi_encoder_trained.pth"

if not os.path.exists(MODEL_PATH):
    print("Model not found. Please ensure you are in the correct directory.")
    exit()

model = MSI_Structural_Encoder()
model.features.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
model.eval()

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

def get_vector(mz_target):
    # Find the file closest to mz_target
    files = [f for f in os.listdir(IMG_DIR) if f.endswith('.png')]
    mzs = [float(f.replace("ion_", "").replace(".png", "")) for f in files]
    idx = np.argmin(np.abs(np.array(mzs) - mz_target))
    closest_mz = mzs[idx]
    
    if abs(closest_mz - mz_target) > 0.1:
        return None, closest_mz
        
    img_path = os.path.join(IMG_DIR, files[idx])
    img = Image.open(img_path).convert("L")
    tensor = transform(img).unsqueeze(0)
    
    with torch.no_grad():
        feat = model(tensor)
        feat_norm = F.normalize(feat, p=2, dim=1)
    return feat_norm.numpy().flatten(), closest_mz

print("--- INTERNAL VALIDATION: ISOTOPIC & STRUCTURAL COHERENCE ---")

# Test 1: Isotopic Coherence (Atropine and its +1 Isotope)
v1, mz1 = get_vector(290.02)
v2, mz2 = get_vector(291.00)
if v1 is not None and v2 is not None:
    sim = np.dot(v1, v2)
    print(f"Isotopic Match -> m/z {mz1:.2f} (Atropine) vs m/z {mz2:.2f} (+1 Isotope): Cosine Similarity = {sim:.4f}")

# Test 2: Structural Coherence (Atropine and Scopolamine - both vascular)
v3, mz3 = get_vector(304.00)
if v1 is not None and v3 is not None:
    sim = np.dot(v1, v3)
    print(f"Biological Match -> m/z {mz1:.2f} (Atropine) vs m/z {mz3:.2f} (Scopolamine): Cosine Similarity = {sim:.4f}")

# Test 3: Unrelated Coherence (Atropine vs random background or non-vascular)
v4, mz4 = get_vector(142.09) # Tropine (paravascular, should be lower than scopolamine)
if v1 is not None and v4 is not None:
    sim = np.dot(v1, v4)
    print(f"Divergent Match -> m/z {mz1:.2f} (Atropine) vs m/z {mz4:.2f} (Tropine): Cosine Similarity = {sim:.4f}")