backend/model_wrapper.py

import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import joblib
import os

# =============================================================================
# EXACT 121 features expected by the model (ALPHABETICALLY SORTED)
# This order matches the non_output_cols minus starttime, endtime, subject_id, row_id
# =============================================================================
MODEL_INPUT_FEATURES = [
    'age', 'albumin_max', 'albumin_min', 'alp_max', 'alp_min', 'alt_max', 'alt_min',
    'aniongap_max', 'aniongap_min', 'antibiotic_count', 'ast_max', 'ast_min',
    'bands_max', 'bands_min', 'baseexcess_max', 'baseexcess_min', 'bicarbonate_max',
    'bicarbonate_min', 'bilirubin_direct_max', 'bilirubin_direct_min',
    'bilirubin_indirect_max', 'bilirubin_indirect_min', 'bilirubin_total_max',
    'bilirubin_total_min', 'bun_max', 'bun_min', 'calcium_max', 'calcium_min',
    'chloride_max', 'chloride_min', 'ck_mb_max', 'ck_mb_min', 'creatinine_max',
    'creatinine_min', 'crp_max', 'crp_min', 'dbp_max', 'dbp_min', 'fibrinogen_max',
    'fibrinogen_min', 'fio2_max', 'fio2_min', 'gcs_eyes_max', 'gcs_eyes_min',
    'gcs_max', 'gcs_min', 'gcs_motor_max', 'gcs_motor_min', 'gcs_verbal_max',
    'gcs_verbal_min', 'gender', 'ggt_max', 'ggt_min', 'globulin_max', 'globulin_min',
    'glucose_max', 'glucose_min', 'heart_rate_max', 'heart_rate_min', 'height',
    'hemoglobin_max', 'hemoglobin_min', 'hr', 'immature_granulocytes_max',
    'immature_granulocytes_min', 'inr_max', 'inr_min', 'lactate_max', 'lactate_min',
    'lymphocytes_abs_max', 'lymphocytes_abs_min', 'mbp_max', 'mbp_min',
    'neutrophils_abs_max', 'neutrophils_abs_min', 'ntprobnp_max', 'ntprobnp_min',
    'pco2_max', 'pco2_min', 'pfratio_max', 'pfratio_min', 'ph_max', 'ph_min',
    'platelet_max', 'platelet_min', 'po2_max', 'po2_min', 'potassium_max',
    'potassium_min', 'pt_max', 'pt_min', 'resp_rate_max', 'resp_rate_min',
    'sbp_max', 'sbp_min', 'so2_max', 'so2_min', 'sodium_max', 'sodium_min',
    'spo2_max', 'spo2_min', 'stay_id', 'temperature_max', 'temperature_min',
    'total_protein_max', 'total_protein_min', 'totalco2_max', 'totalco2_min',
    'troponin_t_max', 'troponin_t_min', 'urineoutput_max', 'urineoutput_min',
    'vaso_dopamine_max', 'vaso_epinephrine_max', 'vaso_norepinephrine_max',
    'vaso_phenylephrine_max', 'vaso_vasopressin_max', 'ventilation_flag',
    'wbc_max', 'wbc_min', 'weight'
]

# --- Model Definitions (Copied from Notebook) ---

class TemporalAttnPool(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.score = nn.Linear(d_model, 1)

    def forward(self, z, padding_mask):
        scores = self.score(z).squeeze(-1)  # [B, T]
        scores = scores.masked_fill(~padding_mask, -1e9)
        alpha = torch.softmax(scores, dim=1)
        pooled = (z * alpha.unsqueeze(-1)).sum(dim=1)
        return pooled


class GRUDTransformer(nn.Module):
    """
    GRU-D Transformer model with multi-head outputs for different window sizes.
    Matches the notebook training architecture exactly.
    """
    def __init__(
        self,
        n_features,
        hidden_size=64,
        d_model=128,
        nhead=4,
        num_layers=2,
        reg_dim=8,
        bin_dim=1
    ):
        super().__init__()

        self.input_size = n_features * 3  # x + mask + delta
        self.gru = nn.GRU(self.input_size, hidden_size, batch_first=True)
        self.to_dmodel = nn.Linear(hidden_size, d_model)

        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.attn_pool = TemporalAttnPool(d_model)

        # 6-head setup: 3 regression heads + 3 binary heads (per window size 6, 12, 24)
        self.reg_heads = nn.ModuleList([nn.Linear(d_model, reg_dim) for _ in range(3)])
        self.bin_heads = nn.ModuleList([nn.Linear(d_model, bin_dim) for _ in range(3)])
        self.heads = nn.ModuleList(list(self.reg_heads) + list(self.bin_heads))

    def forward(self, x, mask, delta, window_id=None):
        inp = torch.cat([x, mask, delta], dim=-1)
        h, _ = self.gru(inp)
        z = self.to_dmodel(h)

        time_mask = mask.sum(dim=-1) > 0
        z = self.transformer(z, src_key_padding_mask=~time_mask)
        pooled = self.attn_pool(z, padding_mask=time_mask)

        # For inference, use window_id=0 (6-hour window) by default
        if window_id is None:
            window_id = torch.zeros(x.size(0), dtype=torch.long, device=x.device)

        y_reg_out = torch.zeros(x.size(0), self.reg_heads[0].out_features, device=x.device)
        y_bin_out = torch.zeros(x.size(0), self.bin_heads[0].out_features, device=x.device)

        for i, w_id in enumerate(window_id):
            y_reg_out[i] = self.reg_heads[w_id](pooled[i])
            y_bin_out[i] = self.bin_heads[w_id](pooled[i])

        return y_reg_out, y_bin_out


# --- Wrapper Class ---

class ModelWrapper:
    def __init__(self, model_dir):
        self.device = torch.device("cpu")
        print(f"Loading model artifacts from {model_dir}...")
        
        # Load scalers and global mean from new model files
        self.scaler_X = joblib.load(os.path.join(model_dir, "scaler_X.pkl"))
        self.scaler_y_reg = joblib.load(os.path.join(model_dir, "scaler_y_reg.pkl"))
        self.global_feat_mean = np.load(os.path.join(model_dir, "global_feat_mean30.npy"))
        
        # Model parameters
        self.n_features = self.scaler_X.mean_.shape[0] if hasattr(self.scaler_X, "mean_") else 121
        self.reg_dim = 8  # respiration, coagulation, liver, cardiovascular, cns, renal, hours_beforesepsis, hours_beforedeath
        self.bin_dim = 1  # sepsis (binary)
        
        print(f"Initializing model with n_features={self.n_features}, reg_dim={self.reg_dim}, bin_dim={self.bin_dim}")
        
        # Initialize model architecture
        self.model = GRUDTransformer(
            n_features=self.n_features,
            hidden_size=64,
            d_model=128,
            nhead=4,
            num_layers=2,
            reg_dim=self.reg_dim,
            bin_dim=self.bin_dim
        )
        
        # Load model weights - try different formats
        model_path = os.path.join(model_dir, "model_joblib.pkl")
        weights_loaded = False
        
        # Custom unpickler to handle CUDA tensors on CPU-only machines
        import pickle
        import io
        
        class CPU_Unpickler(pickle.Unpickler):
            def find_class(self, module, name):
                if module == 'torch.storage' and name == '_load_from_bytes':
                    return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
                else:
                    return super().find_class(module, name)
        
        # Try CPU_Unpickler first (handles CUDA->CPU mapping)
        try:
            with open(model_path, 'rb') as f:
                state_dict = CPU_Unpickler(f).load()
            if isinstance(state_dict, dict) and 'model' in state_dict:
                self.model.load_state_dict(state_dict['model'])
            elif isinstance(state_dict, dict) and 'model_state_dict' in state_dict:
                self.model.load_state_dict(state_dict['model_state_dict'])
            else:
                self.model.load_state_dict(state_dict)
            print("Model weights loaded successfully ✅")
            weights_loaded = True
        except Exception as e:
            print(f"CPU_Unpickler failed: {e}")
            
        # Try torch.load as fallback
        if not weights_loaded:
            try:
                state_dict = torch.load(model_path, map_location='cpu')
                if isinstance(state_dict, dict) and 'model' in state_dict:
                    self.model.load_state_dict(state_dict['model'])
                elif isinstance(state_dict, dict) and 'model_state_dict' in state_dict:
                    self.model.load_state_dict(state_dict['model_state_dict'])
                else:
                    self.model.load_state_dict(state_dict)
                print("Model weights loaded via torch.load ✅")
                weights_loaded = True
            except Exception as e:
                print(f"torch.load failed: {e}")
             
        self.model.to(self.device)
        self.model.eval()

        # Output columns - regression outputs (scaled)
        self.regression_cols = [
            "respiration", "coagulation", "liver", "cardiovascular",
            "cns", "renal", "hours_beforesepsis", "hours_beforedeath"
        ]
        # Binary output (logit -> sigmoid)
        self.binary_cols = ["sepsis"]

    def preprocess_sequence(self, records: list):
        """
        Preprocess patient records for model input with GRU-D style imputation.
        """
        if not records:
            return None, None, None
            
        df = pd.DataFrame(records)
        
        # Map gender (f0_) to numeric
        if 'f0_' in df.columns:
            df['gender'] = df['f0_'].map({'M': 0, 'F': 1, 'Male': 0, 'Female': 1})
            df['gender'] = df['gender'].fillna(0)
        elif 'gender' not in df.columns:
            df['gender'] = np.nan
            
        # Add missing columns
        if 'weight' not in df.columns:
            df['weight'] = np.nan
            
        # Nullify target columns (prevent data leakage)
        target_cols = [
            "respiration", "coagulation", "liver", "cardiovascular",
            "cns", "renal", "hours_beforesepsis", "sepsis",
            "fod", "hours_beforedeath"
        ]
        for col in target_cols:
            if col in df.columns:
                df[col] = np.nan
        
        # Ensure all required features exist
        for col in MODEL_INPUT_FEATURES:
            if col not in df.columns:
                df[col] = np.nan
                
        X_df = df[MODEL_INPUT_FEATURES]
        X_seq = X_df.apply(pd.to_numeric, errors='coerce').to_numpy(dtype=np.float32)
        
        T, F = X_seq.shape
        
        if F != self.n_features:
            print(f"WARNING: Feature count mismatch. Expected {self.n_features}, got {F}.")
            
        # GRU-D style imputation
        mask = ~np.isnan(X_seq)
        X_filled = np.zeros_like(X_seq)
        delta = np.zeros_like(X_seq)
        
        # Get time column (hr) for delta calculation
        if 'hr' in df.columns:
            times = df['hr'].values.astype(float)
        else:
            times = np.arange(T, dtype=float)
            
        for f in range(F):
            mean_val = self.global_feat_mean[f] if f < len(self.global_feat_mean) else 0.0
            last_val = mean_val
            last_time = times[0] if len(times) > 0 else 0
            
            for t in range(T):
                if mask[t, f]:
                    delta[t, f] = 0.0
                    last_val = X_seq[t, f]
                    last_time = times[t]
                    X_filled[t, f] = last_val
                else:
                    if t > 0:
                        delta[t, f] = times[t] - last_time
                    else:
                        delta[t, f] = 0.0
                        
                    gamma = np.exp(-delta[t, f])
                    X_filled[t, f] = gamma * last_val + (1 - gamma) * mean_val
                    last_val = X_filled[t, f]
                    
        # Scale features
        X_scaled = self.scaler_X.transform(X_filled)
        
        # Handle NaN/Inf from scaling (zero-variance features produce NaN)
        X_scaled = np.nan_to_num(X_scaled, nan=0.0, posinf=0.0, neginf=0.0)
        
        # Convert to tensors [1, T, F]
        X_tensor = torch.tensor(X_scaled, dtype=torch.float32).unsqueeze(0)
        mask_tensor = torch.tensor(mask.astype(float), dtype=torch.float32).unsqueeze(0)
        delta_tensor = torch.tensor(delta, dtype=torch.float32).unsqueeze(0)
        
        return X_tensor, mask_tensor, delta_tensor

    def predict(self, records: list, window_id: int = 0):
        """
        Run prediction and return all 10 outputs.
        
        Args:
            records: List of patient records
            window_id: Prediction window (0=6h, 1=12h, 2=24h)
        """
        if not records:
            return None
        
        # Validate window_id
        window_id = max(0, min(2, window_id))  # Clamp to [0, 1, 2]
            
        with torch.no_grad():
            X, mask, delta = self.preprocess_sequence(records)
            if X is None:
                return None
                
            X = X.to(self.device)
            mask = mask.to(self.device)
            delta = delta.to(self.device)
            
            # Use specified window_id (0=6h, 1=12h, 2=24h)
            window_tensor = torch.tensor([window_id], dtype=torch.long, device=self.device)
            
            y_reg_out, y_bin_out = self.model(X, mask, delta, window_tensor)
            
            # Inverse transform regression outputs
            y_reg_np = y_reg_out.cpu().numpy()
            y_reg_original = self.scaler_y_reg.inverse_transform(y_reg_np)
            
            # Apply sigmoid to binary output (trained with BCEWithLogitsLoss)
            y_bin_np = torch.sigmoid(y_bin_out).cpu().numpy()
            
            result = {}
            
            # Regression outputs
            for i, col in enumerate(self.regression_cols):
                val = float(y_reg_original[0, i])
                # Clip SOFA scores to valid range [0, 4]
                if col in ["respiration", "coagulation", "liver", "cardiovascular", "cns", "renal"]:
                    val = max(0.0, min(4.0, val))
                # Clip hours to non-negative
                elif col in ["hours_beforesepsis", "hours_beforedeath"]:
                    val = max(0.0, val)
                result[col] = val
            
            # Binary output (sepsis probability)
            result["sepsis"] = float(y_bin_np[0, 0])
            
            # FOD (failure of organ dysfunction) - calculate from SOFA
            # High SOFA total indicates higher mortality risk
            sofa_sum = sum([
                result.get("respiration", 0),
                result.get("coagulation", 0),
                result.get("liver", 0),
                result.get("cardiovascular", 0),
                result.get("cns", 0),
                result.get("renal", 0)
            ])
            # Map SOFA to mortality probability using sigmoid
            # SOFA >= 11 has ~50% mortality in studies
            result["fod"] = 1.0 / (1.0 + np.exp(-0.3 * (sofa_sum - 8)))
            
            return result