import os
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import numpy as np
import pandas as pd
import torch.nn.functional as F
import matplotlib.pyplot as plt

# 1. CONFIGURATION
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")
OUTPUT_DIR = os.path.join(BASE_DIR, "..", "figures", "guided_validation")
os.makedirs(OUTPUT_DIR, exist_ok=True)

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# FIX: Restore the full list of verified vascular landmarks to build a robust structural vector
ANCHOR_MZS = [155.75, 156.07, 156.40, 273.77, 274.09, 275.07, 286.12, 304.00, 304.33]
SIMI_THRESHOLD = 0.85  # Calibrated to capture true negative space co-exclusions down to 0.85

class MSI_Structural_Encoder(nn.Module):
    """
    The exact high-fidelity custom architecture used during your successful 
    zero-padded 256x256 contrastive training run.
    """
    def __init__(self):
        super(MSI_Structural_Encoder, self).__init__()
        
        self.features = nn.Sequential(
            # Layer 1: Input 256x256x1 (Standardized size)
            nn.Conv2d(1, 32, kernel_size=5, stride=2, padding=2, bias=False),
            nn.GroupNorm(4, 32),
            nn.ReLU(),
            
            # Layer 2
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1, bias=False),
            nn.GroupNorm(8, 64),
            nn.ReLU(),
            
            # Layer 3
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=False),
            nn.GroupNorm(16, 128),
            nn.ReLU(),
            
            # Layer 4
            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()
    # Load only the feature extractor weights to match the save format in CNN_proof_1.py
    model.features.load_state_dict(torch.load(path, map_location=DEVICE))
    print(f"Successfully synchronized weights with custom encoder architecture: {path}")
        
    model.to(DEVICE)
    model.eval()
    return model

if __name__ == "__main__":
    encoder = load_trained_encoder(ENCODER_PATH)
    
    transform = transforms.Compose([
        transforms.Resize((256, 256)), # Standardized size matching training
        transforms.ToTensor()
    ])
    
    # 3. SCAN ALL IMAGE MAPS
    all_files = sorted([f for f in os.listdir(IMG_DIR) if f.endswith('.png')])
    print(f"Discovered {len(all_files)} total ion slices in data pool.")
    
    ion_registry = []
    for f in all_files:
        try:
            mz_extracted = float(f.replace("ion_", "").replace(".png", ""))
            ion_registry.append({"filename": f, "mz": mz_extracted})
        except ValueError:
            continue
            
    df_ions = pd.DataFrame(ion_registry)
    
    # 4. EXTRACT VECTORS
    print("Extracting L2-normalized latent fingerprints...")
    latent_vectors = {}
    
    with torch.no_grad():
        for _, row in df_ions.iterrows():
            fpath = os.path.join(IMG_DIR, row['filename'])
            img_tensor = transform(Image.open(fpath).convert("L")).unsqueeze(0).to(DEVICE)
            feat = encoder(img_tensor)
            feat_norm = F.normalize(feat, p=2, dim=1) # Project onto unit hypersphere
            latent_vectors[row['filename']] = feat_norm.squeeze(0).cpu().numpy()

    # 5. CONSTRUCT SEED
    anchor_features = []
    found_anchors = []
    
    for anchor in ANCHOR_MZS:
        match = df_ions[np.isclose(df_ions['mz'], anchor, atol=0.05)]
        if not match.empty:
            fname = match.iloc[0]['filename']
            anchor_features.append(latent_vectors[fname])
            found_anchors.append(fname)
            
    if len(anchor_features) == 0:
        print("Error: None of your specified anchor vein images were located.")
        exit()
        
    vascular_profile_centroid = np.mean(anchor_features, axis=0)
    vascular_profile_centroid /= np.linalg.norm(vascular_profile_centroid)
    
    print(f"Generated unified vascular profile from {len(found_anchors)} anchor ions.")

    # 6. EVALUATE SIMILARITY
    guided_matches = []
    for fname, vector in latent_vectors.items():
        similarity = np.dot(vector, vascular_profile_centroid)
        if similarity >= SIMI_THRESHOLD:
            mz_val = df_ions[df_ions['filename'] == fname]['mz'].values[0]
            guided_matches.append({"filename": fname, "mz": mz_val, "similarity": similarity})
            
    df_guided = pd.DataFrame(guided_matches).sort_values(by="similarity", ascending=False)
    print(f"\nFound {len(df_guided)} ions displaying vascular alignment (Cosine Sim >= {SIMI_THRESHOLD})")
    
    csv_out = os.path.join(OUTPUT_DIR, "guided_vascular_ions.csv")
    df_guided.to_csv(csv_out, index=False)

    # 7. GENERATE GALLERIES
    max_per_page = 12
    chunks = [df_guided[i:i + max_per_page] for i in range(0, len(df_guided), max_per_page)]
    
    for page_idx, chunk in enumerate(chunks):
        num_plots = len(chunk)
        rows = int(np.ceil(num_plots / 4))
        cols = min(num_plots, 4)
        
        fig, axes = plt.subplots(rows, cols, figsize=(16, rows * 4))
        fig.suptitle(f"Guided Vascular Affinity Grouping - Tier {page_idx + 1}", fontsize=16, y=0.98)
        
        flat_axes = axes.flatten() if isinstance(axes, np.ndarray) else [axes]
        
        for idx, (_, row) in enumerate(chunk.iterrows()):
            ax = flat_axes[idx]
            img_path = os.path.join(IMG_DIR, row['filename'])
            img = Image.open(img_path)
            # Use 'viridis' for improved structural visibility
            ax.imshow(img, cmap='viridis')
            ax.set_title(f"m/z {row['mz']:.2f}\nSim: {row['similarity']:.4f}", fontsize=10)
            ax.axis('off')
            
        for empty_idx in range(num_plots, len(flat_axes)):
            flat_axes[empty_idx].axis('off')
            
        plt.tight_layout()
        gallery_out = os.path.join(OUTPUT_DIR, f"Guided_Vascular_Gallery_Page_{page_idx + 1}.png")
        plt.savefig(gallery_out, dpi=200, bbox_inches='tight')
        plt.close()

    print("\nGuided extraction validation cycle complete.")

    print("\n--- RUNNING ISOTOPIC FUNCTIONAL VALIDATION ---")
    verified_biological_count = 0

    for idx, row in df_guided.iterrows():
        target_mz = row['mz']
        isotope_mz = target_mz + 1.0033

        # Check if the expected isotope exists within your parsed data pool
        isotope_match = df_ions[np.isclose(df_ions['mz'], isotope_mz, atol=0.02)]

        if not isotope_match.empty:
            iso_fname = isotope_match.iloc[0]['filename']
            iso_vector = latent_vectors[iso_fname]

            # Calculate similarity of the isotope to your core vascular vector
            iso_sim = np.dot(iso_vector, vascular_profile_centroid)

            if iso_sim >= 0.80:
                print(f"Verified Structural Pair: Base m/z {target_mz:.2f} (Sim: {row['similarity']:.4f}) -> Isotope m/z {isotope_mz:.2f} (Sim: {iso_sim:.4f})")
                verified_biological_count += 1

    print(f"Validation complete. Confirmed {verified_biological_count} true isotopic structural pairs.")