import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import models, transforms
from PIL import Image
import torch.nn.functional as F
import numpy as np
import random
import time
import psutil

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")
MODEL_SAVE_PATH = os.path.join(BASE_DIR, "..", "models", "msi_encoder_trained.pth")

# VRAM OPTIMIZATION ENGINE
BATCH_SIZE = 8            # Physical batch size (keeps VRAM safely under 3.5 GB)
ACCUMULATION_STEPS = 4    # 8 x 4 = Effective Batch Size of 32!
EPOCHS = 60               # Slightly extended to give the larger batch size room to converge
LEARNING_RATE = 0.0003    # Tuned for an effective batch size of 32
TEMPERATURE = 0.5
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class MSIAugmentations:
    def __init__(self, size=256):  # Enforce 256 dimension consistency
        self.spatial_transform = transforms.Compose([
            transforms.Resize((size, size)),
            transforms.ToTensor()
        ])

    def apply_msi_chemical_noise(self, tensor):
        x = tensor.clone()
        
        # 1. Intensity Scaling (Global Concentration Shifts)
        scale_factor = random.uniform(0.8, 1.2)
        x = x * scale_factor
        
        # 2. Poisson Noise Simulation (Shot Noise)
        # Scaled to avoid breaking the 0.0 - 1.0 float bounds
        vals = len(torch.unique(x))
        vals = 2 ** np.ceil(np.log2(vals)) if vals > 0 else 32
        # Add soft, random Gaussian-Poisson approximation
        noise = torch.randn_like(x) * (torch.sqrt(x + 1e-5) / float(vals))
        x = x + noise

        # 3. Random Missing Values (Standard Dropout - e.g., 2% pixels lost)
        random_dropout_mask = (torch.rand_like(x) > 0.02).float()
        x = x * random_dropout_mask

        # 4. Intensity-Dependent Missing Values (Low-Signal Clipping)
        # Pixels with very low values have a higher chance of dropping out completely
        low_signal_mask = torch.rand_like(x)
        # Drop threshold scales inversely with intensity: lower values drop more easily
        drop_threshold = 0.15 * (1.0 - x) 
        x = torch.where(low_signal_mask > drop_threshold, x, torch.zeros_like(x))

        return torch.clamp(x, 0.0, 1.0)

    def __call__(self, x):
        base_tensor = self.spatial_transform(x)
        # Create two distinct chemical representations of the same underlying coordinate geometry
        return self.apply_msi_chemical_noise(base_tensor), self.apply_msi_chemical_noise(base_tensor)

class SimCLRMSIDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_dir = img_dir
        self.img_names = [f for f in os.listdir(img_dir) if f.endswith('.png')]
        self.transform = transform

    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")
        if self.transform:
            xi, xj = self.transform(image)
        return xi, xj

class MSI_Structural_Encoder(nn.Module):
    def __init__(self, projection_dim=128):
        super(MSI_Structural_Encoder, self).__init__()
        
        # High-Fidelity Custom Encoder for Padded Data
        self.features = nn.Sequential(
            # Input: 256x256x1 (Resize in transform)
            nn.Conv2d(1, 32, kernel_size=5, stride=2, padding=2, bias=False),
            nn.GroupNorm(4, 32), # Replaces BatchNorm for perfect stability at batch size = 8
            nn.ReLU(),
            
            # 112x112 -> 56x56
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1, bias=False),
            nn.GroupNorm(8, 64),
            nn.ReLU(),
            
            # 56x56 -> 28x28 (Captures true raw pixel leaf structures)
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=False),
            nn.GroupNorm(16, 128),
            nn.ReLU(),
            
            # 28x28 -> 14x14
            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1, bias=False),
            nn.GroupNorm(32, 256),
            nn.ReLU(),
            
            nn.AdaptiveAvgPool2d((1, 1)) # Yields a clean 256-dimensional vector
        )
        
        self.projector = nn.Sequential(
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, projection_dim)
        )

    def forward(self, x):
        h = self.features(x).squeeze(-1).squeeze(-1)
        h_norm = F.normalize(h, p=2, dim=1) # Projects onto the unit hypersphere
        z = self.projector(h_norm)
        return h, z

def nt_xent_loss(z1, z2, temperature=0.5):
    batch_size = z1.shape[0]
    z = torch.cat([z1, z2], dim=0)
    z = F.normalize(z, p=2, dim=1)
    
    sim_matrix = torch.mm(z, z.t()) / temperature
    mask = torch.eye(2 * batch_size, device=z.device).bool()
    sim_matrix = sim_matrix.masked_fill(mask, -9e15)
    
    targets = torch.arange(2 * batch_size, device=z.device)
    targets[:batch_size] += batch_size
    targets[batch_size:] -= batch_size
    
    return F.cross_entropy(sim_matrix, targets)

if __name__ == "__main__":
    print(f"Launching VRAM-Optimized Contrastive Pipeline on: {DEVICE}")
    print(f"Physical Batch Size: {BATCH_SIZE} | Accumulation Steps: {ACCUMULATION_STEPS} | Effective Batch Size: {BATCH_SIZE * ACCUMULATION_STEPS}")
    os.makedirs(os.path.dirname(MODEL_SAVE_PATH), exist_ok=True)
    
    aug_pipeline = MSIAugmentations(size=256)
    dataset = SimCLRMSIDataset(IMG_DIR, transform=aug_pipeline)
    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True, num_workers=2, pin_memory=True)

    model = MSI_Structural_Encoder().to(DEVICE)
    optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS, eta_min=1e-6)

    process = psutil.Process(os.getpid())
    
    # Environment tweak to combat memory fragmentation
    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
    torch.cuda.empty_cache()

    for epoch in range(EPOCHS):
        model.train()
        total_loss = 0
        epoch_start = time.time()
        
        optimizer.zero_grad() # Initialize gradients outside the accumulation loop
        
        for batch_idx, (xi, xj) in enumerate(loader):
            xi, xj = xi.to(DEVICE), xj.to(DEVICE)
            
            _, z1 = model(xi)
            _, z2 = model(xj)
            
            # Scale loss by accumulation steps to ensure correct gradient weighting
            loss = nt_xent_loss(z1, z2, temperature=TEMPERATURE) / ACCUMULATION_STEPS
            loss.backward()
            
            total_loss += loss.item() * ACCUMULATION_STEPS
            
            # Step the optimizer only after gathering enough gradients
            if (batch_idx + 1) % ACCUMULATION_STEPS == 0 or (batch_idx + 1) == len(loader):
                optimizer.step()
                optimizer.zero_grad()
        
        scheduler.step()
        
        # Periodic cache flushing keeps VRAM overhead low
        if epoch % 5 == 0:
            torch.cuda.empty_cache()
            
        print(f"Epoch [{epoch+1}/{EPOCHS}], Loss: {total_loss / len(loader):.4f}, Time: {time.time()-epoch_start:.2f}s, RAM Used: {process.memory_info().rss / 1024 / 1024:.1f} MB")

    print(f"--> Exporting optimized 32-BS encoder weights to: {MODEL_SAVE_PATH}")
    torch.save(model.features.state_dict(), MODEL_SAVE_PATH)
    print(f"Model saved to: {MODEL_SAVE_PATH}")