import torch
import torch.nn as nn
import os

# ==========================================
# 1. MODEL CONFIGURATION (V7 Standard)
# ==========================================
# Fixed latent dimension matching your trained "amlm_v5.pth"
LATENT_DIM = 64

class AMLM_3D(nn.Module):
    def __init__(self):
        super(AMLM_3D, self).__init__()
        
        # --- ENCODER ---
        # Input: (Batch, 6, 64, 9) 
        # 6 Channels = 3 Value Channels (Data) + 3 Visibility Channels (Mask)
        self.encoder = nn.Sequential(
            # Layer 1: 64x9 -> 32x9
            nn.Conv2d(6, 32, kernel_size=3, padding=1), 
            nn.ReLU(), 
            nn.MaxPool2d((2, 1)), 
            
            # Layer 2: 32x9 -> 16x9
            nn.Conv2d(32, 64, kernel_size=3, padding=1), 
            nn.ReLU(), 
            nn.MaxPool2d((2, 1)), 
            
            # Flatten: 64 channels * 16 height * 9 width = 9216
            nn.Flatten() 
        )
        
        # --- BOTTLENECK ---
        # Compresses the 9216 features into the 64-dim "Clinical Syndrome" vector
        self.bottleneck = nn.Sequential(
            nn.Linear(9216, 256), 
            nn.ReLU(), 
            nn.Linear(256, LATENT_DIM) 
        )
        
        # --- DECODER INPUT ---
        # Expands 64-dim syndrome back to spatial features
        self.decoder_input = nn.Sequential(
            nn.Linear(LATENT_DIM, 256), 
            nn.ReLU(), 
            nn.Linear(256, 9216), 
            nn.ReLU()
        )
        
        # --- DECODER ---
        # Reconstructs the full probability map
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 32, kernel_size=3, padding=1), 
            nn.Upsample(scale_factor=(2, 1)), 
            nn.ReLU(),
            
            nn.ConvTranspose2d(32, 16, kernel_size=3, padding=1), 
            nn.Upsample(scale_factor=(2, 1)), 
            nn.ReLU(),
            
            # Output: 3 Channels (Initial, Medial, Final)
            nn.Conv2d(16, 3, kernel_size=3, padding=1), 
            nn.Sigmoid() # Squish output to 0.0 - 1.0 probability
        )

    def forward(self, x):
        """
        Forward pass for inference.
        x shape must be: (Batch_Size, 6, 64, 9)
        """
        features = self.encoder(x)
        latent = self.bottleneck(features) 
        
        x_recon = self.decoder_input(latent)
        # Reshape back to spatial dimensions (Batch, Channels, Height, Width)
        x_recon = x_recon.view(-1, 64, 16, 9) 
        
        # Output map (Batch, 3, 64, 9)
        output = self.decoder(x_recon)
        
        return output, latent

# ==========================================
# 2. LOADER UTILITY
# ==========================================
def load_amlm_model(model_path, device):
    """
    Initializes the model and loads weights safely.
    """
    print(f"--- Loading AMLM V7 from {model_path} ---")
    model = AMLM_3D().to(device)
    
    if os.path.exists(model_path):
        try:
            state_dict = torch.load(model_path, map_location=device)
            model.load_state_dict(state_dict)
            model.eval() # Set to inference mode (freezes Dropout/BatchNorm)
            print("✅ AMLM Weights Loaded Successfully")
        except Exception as e:
            print(f"❌ CRITICAL: Failed to load AMLM weights: {e}")
            # We assume the app should crash or handle this if weights are bad
            raise e
    else:
        print(f"⚠️ WARNING: Model file not found at {model_path}. Running with random weights (Debugging Mode).")
    
    return model