architectures.py

import logging
import torch
import torch.nn as nn
import torch.nn.functional as F
import json
import time
import pytorch_lightning as pl
from torch.optim import AdamW
from torchmetrics import MeanSquaredError, PearsonCorrCoef, SpearmanCorrCoef, R2Score

logger = logging.getLogger(__name__)

# ===================== VERSION STRING FOR CLUSTER VERIFICATION =====================
ARCH_VERSION = "2024-12-24-stability-fix"
print(f"[ARCH] architectures.py loaded: {ARCH_VERSION}")
# ====================================================================================

class Interp1d(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, y, xnew):
        is_flat = {}
        vals = {'x': x, 'y': y, 'xnew': xnew}
        for name, arr in vals.items():
            is_flat[name] = (arr.dim() == 1)
            if is_flat[name]:
                vals[name] = arr.unsqueeze(0)
        x_2d, y_2d, xnew_2d = vals['x'], vals['y'], vals['xnew']
        B, Nx = x_2d.shape

        # SAFETY: Handle edge case where sequence length is < 5
        if Nx < 5:
            # Return constant interpolation (repeat/average the values)
            ynew_2d = y_2d.mean(dim=1, keepdim=True).expand(-1, xnew_2d.shape[1])
            ctx.save_for_backward(x_2d, y_2d, xnew_2d, 
                                  torch.zeros_like(xnew_2d, dtype=torch.long),
                                  torch.zeros_like(xnew_2d))
            ctx.Nx_was_small = True
            if is_flat['x'] and is_flat['xnew']:
                ynew_2d = ynew_2d.squeeze(0)
            return ynew_2d
        
        ctx.Nx_was_small = False
        idx = torch.searchsorted(x_2d, xnew_2d, right=False) - 1
        idx = idx.clamp(min=0, max=Nx-2)

        xL = torch.gather(x_2d, 1, idx)
        xR = torch.gather(x_2d, 1, idx+1)
        yL = torch.gather(y_2d, 1, idx)
        yR = torch.gather(y_2d, 1, idx+1)

        denom = (xR - xL)
        denom[denom == 0] = 1e-12
        t = (xnew_2d - xL)/denom
        ynew_2d = yL + (yR - yL)*t

        ctx.save_for_backward(x_2d, y_2d, xnew_2d, idx, t)
        if is_flat['x'] and is_flat['xnew']:
            ynew_2d = ynew_2d.squeeze(0)
        return ynew_2d

    @staticmethod
    def backward(ctx, grad_out):
        x_2d, y_2d, xnew_2d, idx, t = ctx.saved_tensors
        grad_x = grad_y = grad_xnew = None
        
        # Handle edge case from forward
        if getattr(ctx, 'Nx_was_small', False):
            if ctx.needs_input_grad[1]:
                grad_y = grad_out.sum(dim=-1, keepdim=True).expand_as(y_2d)
            return grad_x, grad_y, grad_xnew
            
        if ctx.needs_input_grad[1]:
            grad_y_tmp = torch.zeros_like(y_2d)
            idxp1 = (idx + 1).clamp(max=y_2d.shape[1] - 1)  # SAFETY: clamp idxp1
            
            # Calculate gradients
            grad_yL = (1.0 - t) * grad_out
            grad_yR = t * grad_out
            
            # Ensure consistent dtype between source and destination tensors
            grad_yL = grad_yL.to(dtype=grad_y_tmp.dtype)
            grad_yR = grad_yR.to(dtype=grad_y_tmp.dtype)
            
            grad_y_tmp.scatter_add_(1, idx, grad_yL)
            grad_y_tmp.scatter_add_(1, idxp1, grad_yR)
            grad_y = grad_y_tmp
        return grad_x, grad_y, grad_xnew

def interp1d(x, y, xnew):
    return Interp1d.apply(x, y, xnew)


class SWE_Pooling(nn.Module):
    """
    Sliced-Wasserstein Embedding (SWE) Pooling.
    Maps token embeddings [B, L, d_in] => [B, num_slices].
    """
    def __init__(self, d_in, num_slices, num_ref_points, freeze_swe=False):
        super().__init__()
        self.num_slices = num_slices
        self.num_ref_points = num_ref_points

        ref = torch.linspace(-1,1,num_ref_points).unsqueeze(1).repeat(1,num_slices)
        self.reference = nn.Parameter(ref, requires_grad=not freeze_swe)

        self.theta = nn.utils.weight_norm(nn.Linear(d_in, num_slices, bias=False), dim=0)
        self.theta.weight_g.data = torch.ones_like(self.theta.weight_g.data)
        self.theta.weight_g.requires_grad=False
        nn.init.normal_(self.theta.weight_v)

        self.weight = nn.Linear(num_ref_points,1,bias=False)

        if freeze_swe:
            self.theta.weight_v.requires_grad=False
            self.reference.requires_grad=False

    def forward(self, X, mask=None):
        B, N, D = X.shape
        device = X.device

        X_slices = self.theta(X)  # => [B,N,num_slices]
        X_slices_sorted, _ = torch.sort(X_slices, dim=1)

        x_coord = torch.linspace(0,1,N,device=device).unsqueeze(0).repeat(B*self.num_slices,1)
        X_flat = X_slices_sorted.permute(0,2,1).reshape(B*self.num_slices, N)
        xnew = torch.linspace(0,1,self.num_ref_points,device=device).unsqueeze(0).repeat(B*self.num_slices,1)

        y_intp = interp1d(x_coord, X_flat, xnew)
        X_slices_sorted_interp = y_intp.view(B,self.num_slices,self.num_ref_points).permute(0,2,1)

        r_expanded = self.reference.expand_as(X_slices_sorted_interp)
        embeddings = (r_expanded - X_slices_sorted_interp).permute(0,2,1)  # => [B,num_slices,num_ref_points]
        weighted = self.weight(embeddings).sum(dim=-1)  # => [B, num_slices]
        return weighted


#############################################################################
#             Enhanced Mutation-Aware SWE_Pooling                           #
#############################################################################

class MutationAwareSWEPooling(nn.Module):
    """
    Enhanced Sliced-Wasserstein Embedding Pooling with explicit mutation position handling.
    Maps token embeddings [B, L, d_in] => [B, num_slices].
    
    - Preserves mutation position information through weighted aggregation
    - Uses mutation positions to guide the pooling process
    """
    def __init__(self, d_in, num_slices, num_ref_points, freeze_swe=False):
        super().__init__()
        self.num_slices = num_slices
        self.num_ref_points = num_ref_points
        self.d_esm = 1152  # FIXED: Hardcode to 1152 to avoid channel indexing bugs with context window

        # Standard SWE components
        ref = torch.linspace(-1, 1, num_ref_points).unsqueeze(1).repeat(1, num_slices)
        self.reference = nn.Parameter(ref, requires_grad=not freeze_swe)

        # For ESM features (without mutation channel)
        self.theta = nn.utils.weight_norm(nn.Linear(self.d_esm, num_slices, bias=False), dim=0)
        self.theta.weight_g.data = torch.ones_like(self.theta.weight_g.data)
        self.theta.weight_g.requires_grad = False
        nn.init.normal_(self.theta.weight_v)

        # Mutation-aware components
        self.mutation_importance = nn.Sequential(
            nn.Linear(1, 32),
            nn.ReLU(),
            nn.Linear(32, num_slices),
            nn.Sigmoid()
        )
        
        # Position-specific weighting for each slice
        self.pos_weighting = nn.Linear(1, num_slices, bias=False)

        # FIXED: Direct projection of mutation channel (the 1153rd dim)
        # Previously, this channel was only used as a multiplier, meaning if ESM 
        # features had no diff, the result was 0. Now we project it directly.
        self.mut_projection = nn.Linear(1, num_slices, bias=False)
        
        # Final weighting
        self.weight = nn.Linear(num_ref_points, 1, bias=False)

        if freeze_swe:
            self.theta.weight_v.requires_grad = False
            self.reference.requires_grad = False

    def forward(self, X, mask=None):
        """
        X: [B, L, d_in] where d_in = d_esm + 1 (mutation channel)
        mask: [B, L] boolean mask
        """
        B, N, D = X.shape
        device = X.device
        
        # Check if using context window (additional channel)
        use_context = (D > self.d_esm + 1)
        
        if use_context:
            # Split ESM features and channels
            X_esm = X[:, :, :-2]  # [B, L, d_esm]
            X_mut = X[:, :, -2:-1]  # [B, L, 1] - mutation indicator
        else:
            # Split ESM features and mutation channel
            X_esm = X[:, :, :-1]  # [B, L, d_esm]
            X_mut = X[:, :, -1:]  # [B, L, 1] - mutation indicator
        
        # Regular SWE on ESM features
        X_slices = self.theta(X_esm)  # => [B, L, num_slices]
        
        # Compute mutation importance weights
        mut_weights = self.mutation_importance(X_mut)  # [B, L, num_slices]
        
        # Create position encodings (0 to 1 for each sequence)
        pos_tensor = torch.linspace(0, 1, N, device=device).view(1, N, 1).expand(B, N, 1)
        pos_weights = self.pos_weighting(pos_tensor)  # [B, L, num_slices]
        
        # Apply mutation-aware weighting to slices
        # Use both mutation indicator and position information
        # BUGFIX: We ALSO add the projected mutation signal directly. 
        # This ensures the model 'sees' the 1.0 signal even if ESM features are identical.
        X_slices = X_slices * (1.0 + mut_weights * pos_weights) + self.mut_projection(X_mut)
        
        # Sort slices as in standard SWE
        X_slices_sorted, _ = torch.sort(X_slices, dim=1)
        
        # Continue with standard SWE interpolation
        x_coord = torch.linspace(0, 1, N, device=device).unsqueeze(0).repeat(B*self.num_slices, 1)
        X_flat = X_slices_sorted.permute(0, 2, 1).reshape(B*self.num_slices, N)
        xnew = torch.linspace(0, 1, self.num_ref_points, device=device).unsqueeze(0).repeat(B*self.num_slices, 1)
        
        y_intp = interp1d(x_coord, X_flat, xnew)
        X_slices_sorted_interp = y_intp.view(B, self.num_slices, self.num_ref_points).permute(0, 2, 1)
        
        r_expanded = self.reference.expand_as(X_slices_sorted_interp)
        embeddings = (r_expanded - X_slices_sorted_interp).permute(0, 2, 1)  # => [B, num_slices, num_ref_points]
        weighted = self.weight(embeddings).sum(dim=-1)  # => [B, num_slices]
        
        return weighted


#############################################################################
#            Mutation-Specific Cross-Attention with Gating                  #
#############################################################################

class MutationSpecificAttention(nn.Module):
    """
    Enhanced cross-attention that explicitly handles mutation positions with gating.
    - Keeps ESM embeddings (1152-dim) and mutation channel separate
    - Uses specific mutation positions to guide attention
    - Preserves position-specific information throughout the network
    - Adds gating mechanism to control information flow
    - Includes memory-efficient computation for long sequences
    """
    def __init__(self, d_model=1152, num_heads=4, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        
        # Core attention for ESM embeddings only (1152-dim)
        self.query = nn.Linear(d_model, d_model)
        self.key = nn.Linear(d_model, d_model)
        self.value = nn.Linear(d_model, d_model)
        
        # Absolute position encoding
        self.pos_encoder = nn.Sequential(
            nn.Linear(1, 32),
            nn.ReLU(),
            nn.Linear(32, d_model)
        )
        
        # Mutation-position specific attention
        self.mut_encoder = nn.Sequential(
            nn.Linear(2, 64),  # Input: [mut_binary, position_normalized]
            nn.ReLU(),
            nn.Linear(64, num_heads)
        )
        
        # Gating mechanism to control information flow
        self.gate = nn.Sequential(
            nn.Linear(d_model*2, d_model),
            nn.Sigmoid()
        )
        
        self.dropout = nn.Dropout(dropout)
        self.out_proj = nn.Linear(d_model, d_model)
        
    def split_heads(self, x):
        """Split the last dimension into (heads, head_dim)"""
        batch_size, seq_len, _ = x.shape
        x = x.view(batch_size, seq_len, self.num_heads, self.head_dim)
        return x.permute(0, 2, 1, 3)  # [batch, heads, seq_len, head_dim]
    
    def merge_heads(self, x):
        """Merge the (heads, head_dim) into d_model"""
        batch_size, _, seq_len, _ = x.shape
        x = x.permute(0, 2, 1, 3)  # [batch, seq_len, heads, head_dim]
        return x.reshape(batch_size, seq_len, self.d_model)
    
    def forward(self, q_esm, k_esm, v_esm, q_mut, k_mut, mask=None):
        """
        Inputs:
            q_esm, k_esm, v_esm: ESM embeddings [B, L, 1152]
            q_mut, k_mut: Mutation information [B, L, 1]
            mask: Optional attention mask [B, L] or [B, 1, L]
        """
        batch_size = q_esm.shape[0]
        q_len, k_len = q_esm.shape[1], k_esm.shape[1]
        
        # Create position tensors (0-1 range for each sequence)
        q_pos = torch.linspace(0, 1, q_len, device=q_esm.device).view(1, -1, 1).expand(batch_size, q_len, 1)
        k_pos = torch.linspace(0, 1, k_len, device=k_esm.device).view(1, -1, 1).expand(batch_size, k_len, 1)
        
        # Position encoding
        q_pos_enc = self.pos_encoder(q_pos)
        k_pos_enc = self.pos_encoder(k_pos)
        
        # Add position encodings to ESM features
        q_esm_pos = q_esm + q_pos_enc
        k_esm_pos = k_esm + k_pos_enc
        
        # Process core ESM embeddings with position information
        q = self.split_heads(self.query(q_esm_pos))  # [B, h, q_len, d_k]
        k = self.split_heads(self.key(k_esm_pos))    # [B, h, k_len, d_k]
        v = self.split_heads(self.value(v_esm))      # [B, h, v_len, d_v]
        
        # Concatenate mutation indicator with position
        q_mut_pos = torch.cat([q_mut, q_pos], dim=-1)  # [B, q_len, 2]
        k_mut_pos = torch.cat([k_mut, k_pos], dim=-1)  # [B, k_len, 2]
        
        # Encode position-aware mutation information
        q_mut_enc = self.mut_encoder(q_mut_pos)  # [B, q_len, num_heads]
        k_mut_enc = self.mut_encoder(k_mut_pos)  # [B, k_len, num_heads]
        
        # Standard scaled dot-product attention
        d_k = q.size(-1)
        scores = torch.matmul(q, k.transpose(-2, -1)) / (d_k ** 0.5)  # [B, h, q_len, k_len]
        
        # Create mutation-position attention bias
        # This explicitly boosts attention between positions based on mutation status
        mut_attn_bias = torch.matmul(
            q_mut_enc.permute(0, 2, 1).unsqueeze(3),  # [B, h, q_len, 1]
            k_mut_enc.permute(0, 2, 1).unsqueeze(2)   # [B, h, 1, k_len]
        )  # [B, h, q_len, k_len]
        
        # Apply mutation bias to attention scores
        # This makes mutations and their surrounding context attend more to each other
        scores = scores + mut_attn_bias
        
        # Apply mask if provided
        if mask is not None:
            # Fix mask dimension to match scores
            # mask shape should be [B, L] or [B, 1, L]
            if mask.dim() == 2:  # [B, L]
                # For keys mask [B, k_len] -> [B, 1, 1, k_len]
                mask = mask.unsqueeze(1).unsqueeze(2)
            elif mask.dim() == 3 and mask.size(1) == 1:  # [B, 1, L]
                # For keys mask [B, 1, k_len] -> [B, 1, 1, k_len]
                mask = mask.unsqueeze(2)
                
            # Expand mask to match scores dimensions
            # [B, 1, 1, k_len] -> [B, h, q_len, k_len]
            mask = mask.expand(-1, scores.size(1), scores.size(2), -1)
            
            # FIXED: Use -1e4 instead of -1e9 to avoid half-precision overflow
            scores = scores.masked_fill(mask == 0, -1e4)
        
        # Apply softmax and dropout
        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)
        
        # Apply attention to values
        context = torch.matmul(attention_weights, v)  # [B, h, q_len, d_v]
        context = self.merge_heads(context)  # [B, q_len, d_model]
        attn_output = self.out_proj(context)
        
        # Apply gating mechanism (new addition)
        # Concatenate the original query with the attention output to determine the gate
        gate_input = torch.cat([q_esm, attn_output], dim=-1)
        gate_value = self.gate(gate_input)
        
        # Memory optimization for long sequences
        # Processing the gating operation in chunks to prevent OOM errors
        if q_len > 1000:  # Only use chunking for very long sequences
            chunk_size = 500
            output_chunks = []
            
            for i in range(0, q_len, chunk_size):
                end_idx = min(i + chunk_size, q_len)
                # Process chunks
                chunk_gate = gate_value[:, i:end_idx, :]
                chunk_attn = attn_output[:, i:end_idx, :]
                chunk_q = q_esm[:, i:end_idx, :]
                
                # Apply gating equation to this chunk
                chunk_output = chunk_gate * chunk_attn + (1 - chunk_gate) * chunk_q
                output_chunks.append(chunk_output)
            
            # Combine chunks
            output = torch.cat(output_chunks, dim=1)
        else:
            # Original operation for shorter sequences
            output = gate_value * attn_output + (1 - gate_value) * q_esm
        
        return output


class MutationSpecificCrossAttentionBlock(nn.Module):
    """
    Cross-attention block with explicit mutation position handling.
    Each block processes ESM embeddings and mutation channels separately,
    with special emphasis on mutation positions.
    """
    def __init__(self, d_model=1152, num_heads=4, ffn_dim=2048, dropout=0.1):
        super().__init__()
        # Mutation-aware cross attention
        self.attn_c12 = MutationSpecificAttention(d_model, num_heads, dropout)
        self.attn_c21 = MutationSpecificAttention(d_model, num_heads, dropout)
        
        # Layer normalization for ESM embeddings
        self.norm_c1 = nn.LayerNorm(d_model)
        self.norm_c2 = nn.LayerNorm(d_model)
        
        # FFN for ESM embeddings
        self.ffn_c1 = nn.Sequential(
            nn.Linear(d_model, ffn_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, d_model)
        )
        self.ffn_c2 = nn.Sequential(
            nn.Linear(d_model, ffn_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, d_model)
        )
        self.norm_ffn_c1 = nn.LayerNorm(d_model)
        self.norm_ffn_c2 = nn.LayerNorm(d_model)
        
        # Mutation importance update layer
        self.mut_update = nn.Sequential(
            nn.Linear(d_model + 1, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )

    def forward(self, c1_esm, c1_mut, c2_esm, c2_mut, mask1=None, mask2=None):
        """
        Inputs:
            c1_esm, c2_esm: ESM embeddings [B, L, 1152]
            c1_mut, c2_mut: Mutation channels [B, L, 1]
            mask1, mask2: Optional masks
        """
        # c1->c2 cross-attention
        c1_attn = self.attn_c12(c1_esm, c2_esm, c2_esm, c1_mut, c2_mut, mask2)
        c1_out = self.norm_c1(c1_esm + c1_attn)
        
        # c2->c1 cross-attention
        c2_attn = self.attn_c21(c2_esm, c1_esm, c1_esm, c2_mut, c1_mut, mask1)
        c2_out = self.norm_c2(c2_esm + c2_attn)
        
        # Feed-forward
        c1_ffn = self.ffn_c1(c1_out)
        c1_ffn_out = self.norm_ffn_c1(c1_out + c1_ffn)
        
        c2_ffn = self.ffn_c2(c2_out)
        c2_ffn_out = self.norm_ffn_c2(c2_out + c2_ffn)
        
        # Update mutation importance based on attention output
        # This creates a feedback loop where mutation effect is refined
        c1_mut_in = torch.cat([c1_ffn_out, c1_mut], dim=-1)
        c2_mut_in = torch.cat([c2_ffn_out, c2_mut], dim=-1)
        
        # Stabilized update: convex combination ensures values stay in [0, 1]
        # Avoids exponential decay (old bug) and unbounded growth (additive bug)
        c1_mut_updated = 0.9 * c1_mut + 0.1 * self.mut_update(c1_mut_in)
        c2_mut_updated = 0.9 * c2_mut + 0.1 * self.mut_update(c2_mut_in)
        
        return c1_ffn_out, c2_ffn_out, c1_mut_updated, c2_mut_updated


class MutationSpecificCrossAttentionStack(nn.Module):
    """
    Stack of Mutation-Specific Cross-Attention blocks.
    Emphasizes mutation positions throughout the network.
    Now includes gradient checkpointing for memory efficiency.
    """
    def __init__(self, d_model=1152, num_heads=4, ffn_dim=2048, dropout=0.1, num_layers=2):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.use_checkpoint = True  # Enable gradient checkpointing by default
        
        self.blocks = nn.ModuleList([
            MutationSpecificCrossAttentionBlock(
                d_model=d_model,
                num_heads=num_heads,
                ffn_dim=ffn_dim,
                dropout=dropout
            ) for _ in range(num_layers)
        ])

    def forward(self, c1, c2, mask1=None, mask2=None):
        """
        Process protein chains with mutation-specific attention.
        c1, c2: [B, L, D] where D can be 1153 (original) or 1154 (with context window)
        Uses gradient checkpointing when in training mode to save memory.
        """
        # Check input dimension to determine if context window is used
        d_in = c1.shape[2]
        use_context = (d_in > 1153)
        
        if use_context:
            # Split ESM embeddings from mutation+context channels
            c1_esm, c1_channels = c1[:, :, :-2], c1[:, :, -2:]  # [B, L, 1152], [B, L, 2]
            c2_esm, c2_channels = c2[:, :, :-2], c2[:, :, -2:]  # [B, L, 1152], [B, L, 2]
            
            # Extract mutation channel (first channel)
            c1_mut = c1_channels[:, :, :1]  # [B, L, 1]
            c2_mut = c2_channels[:, :, :1]  # [B, L, 1]
        else:
            # Original behavior - just split ESM and mutation
            c1_esm, c1_mut = c1[:, :, :-1], c1[:, :, -1:]  # [B, L, 1152], [B, L, 1]
            c2_esm, c2_mut = c2[:, :, :-1], c2[:, :, -1:]  # [B, L, 1152], [B, L, 1]
        
        # Process through attention blocks with optional checkpointing
        for block in self.blocks:
            # Use gradient checkpointing in training mode for memory efficiency
            if self.use_checkpoint and self.training:
                # Define helper function for checkpointing that handles None masks
                def create_checkpoint_fn(block_fn):
                    def checkpoint_fn(esm1, mut1, esm2, mut2, has_mask1, has_mask2, mask1_val, mask2_val):
                        # Conditionally use the masks based on the has_mask flags
                        m1 = mask1_val if has_mask1 else None
                        m2 = mask2_val if has_mask2 else None
                        return block_fn(esm1, mut1, esm2, mut2, m1, m2)
                    return checkpoint_fn
                
                # Convert None masks to flags and dummy tensors for checkpointing
                has_mask1 = mask1 is not None
                has_mask2 = mask2 is not None
                mask1_val = mask1 if has_mask1 else torch.zeros(1, device=c1_esm.device)
                mask2_val = mask2 if has_mask2 else torch.zeros(1, device=c1_esm.device)
                
                # Apply checkpointing
                c1_esm, c2_esm, c1_mut, c2_mut = torch.utils.checkpoint.checkpoint(
                    create_checkpoint_fn(block),
                    c1_esm, c1_mut, c2_esm, c2_mut,
                    torch.tensor(has_mask1, device=c1_esm.device),
                    torch.tensor(has_mask2, device=c1_esm.device),
                    mask1_val, mask2_val
                )
            else:
                c1_esm, c2_esm, c1_mut, c2_mut = block(c1_esm, c1_mut, c2_esm, c2_mut, mask1, mask2)
        
        # Recombine with appropriate channels
        if use_context:
            # Need to preserve the context channel
            context_channels_c1 = c1_channels[:, :, 1:]  # [B, L, 1]
            context_channels_c2 = c2_channels[:, :, 1:]  # [B, L, 1]
            c1_out = torch.cat([c1_esm, c1_mut, context_channels_c1], dim=-1)  # [B, L, 1154]
            c2_out = torch.cat([c2_esm, c2_mut, context_channels_c2], dim=-1)  # [B, L, 1154]
        else:
            # Original behavior
            c1_out = torch.cat([c1_esm, c1_mut], dim=-1)  # [B, L, 1153]
            c2_out = torch.cat([c2_esm, c2_mut], dim=-1)  # [B, L, 1153]
        
        return c1_out, c2_out


#############################################################################
#        AffinityPredictor with Improved Memory Efficiency                  #
#############################################################################

class AffinityPredictor(nn.Module):
    """
    Enhanced AffinityPredictor with explicit mutation position handling.
    embedding_method => "difference", "cosine", "cross_attention", or "cross_attention_swe".
    """
    def __init__(
        self,
        input_dim=1153,  # 1152 (ESM) + 1 (mutation)
        latent_dim=1024,
        num_slices=1024,
        num_ref_points=128,
        dropout_rate=0.2,
        freeze_swe=False,
        embedding_method="difference",
        normalize_difference=False,
        num_hidden_layers=2,
        # cross-attn
        num_cross_attn_layers=2,
        num_attention_heads=4,
        cross_ffn_dim=2048,
    ):
        super().__init__()
        self.embedding_method = embedding_method.lower()
        self.normalize_difference = normalize_difference
        self.input_dim = input_dim
        
        # ESM dimension (without mutation channel)
        self.esm_dim = input_dim - 1  # 1152
        
        # Define cross-attention stack if needed
        self.cross_stack = None
        if "cross_attention" in self.embedding_method:
            self.cross_stack = MutationSpecificCrossAttentionStack(
                d_model=self.esm_dim,  # 1152
                num_heads=num_attention_heads,
                ffn_dim=cross_ffn_dim,
                dropout=dropout_rate,
                num_layers=num_cross_attn_layers
            )
        
        # Enhanced Mutation-Aware SWE Pooling
        self.swe_pooling = None
        if self.embedding_method in ["difference", "cosine", "cross_attention_swe"]:
            # For SWE, we use the full input_dim (1153)
            self.swe_pooling = MutationAwareSWEPooling(
                d_in=input_dim,
                num_slices=num_slices,
                num_ref_points=num_ref_points,
                freeze_swe=freeze_swe
            )
        
        # Define aggregator MLP in-dimensions
        if self.embedding_method == "cosine":
            in_features = 1
        else:
            in_features = num_slices  # difference or cross_attention_swe => [B, num_slices]
        
        # Add projection layer for cross_attention to avoid dynamic creation
        self.cross_attn_projection = None
        if self.embedding_method == "cross_attention":
            cross_proj_in = input_dim  # Full dimension including mutation channel
            cross_proj_out = in_features
            self.cross_attn_projection = nn.Linear(cross_proj_in, cross_proj_out, bias=False)
        
        # Final MLP
        layers = []
        current_dim = in_features
        for _ in range(num_hidden_layers):
            layers.append(nn.Linear(current_dim, latent_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout_rate))
            current_dim = latent_dim
        layers.append(nn.Linear(current_dim, 1))
        self.mlp = nn.Sequential(*layers)

    def forward(self, chain1, chain1_mask, chain2, chain2_mask):
        """
        chain1, chain2 => [B, L, input_dim] (1153 or 1154 with context)
        """
        if "cross_attention" in self.embedding_method:
            # Process through mutation-specific cross-attention
            c1_out, c2_out = self.cross_stack(chain1, chain2, chain1_mask, chain2_mask)
            
            if self.embedding_method == "cross_attention_swe":
                # Apply enhanced SWE pooling
                rep1 = self.swe_pooling(c1_out, chain1_mask)  # [B, num_slices]
                rep2 = self.swe_pooling(c2_out, chain2_mask)  # [B, num_slices]
                
                # Difference aggregator
                diff = rep1 - rep2
                if self.normalize_difference:
                    diff = F.normalize(diff, p=2, dim=1)
                    
                # Final prediction
                preds = self.mlp(diff).squeeze(-1)
                
                return preds
                
            elif self.embedding_method == "cross_attention":
                # Use mutation-weighted pooling
                # Extract mutation channel to guide pooling
                d_in = c1_out.shape[2]
                use_context = (d_in > 1153)
                if use_context:
                    c1_mut = c1_out[:, :, -2:-1]  # [B, L, 1]
                    c2_mut = c2_out[:, :, -2:-1]  # [B, L, 1]
                else:
                    c1_mut = c1_out[:, :, -1:]  # [B, L, 1]
                    c2_mut = c2_out[:, :, -1:]  # [B, L, 1]
                
                # Weighted pooling - gives higher weight to mutated positions
                c1_weights = F.softmax(c1_mut * 10, dim=1)  # Sharpen weights
                c2_weights = F.softmax(c2_mut * 10, dim=1)
                
                c1_pool = torch.sum(c1_out * c1_weights, dim=1)  # [B, 1153/1154]
                c2_pool = torch.sum(c2_out * c2_weights, dim=1)  # [B, 1153/1154]
                
                # Create difference representation
                diff = c1_pool - c2_pool
                if self.normalize_difference:
                    diff = F.normalize(diff, p=2, dim=1)
                
                # Use pre-defined projection layer instead of creating one dynamically
                if self.cross_attn_projection is not None:
                    diff = self.cross_attn_projection(diff)
                
                preds = self.mlp(diff).squeeze(-1)
                
                return preds
        
        elif self.embedding_method == "cosine":
            # Enhanced SWE => [B, num_slices]
            rep1 = self.swe_pooling(chain1, chain1_mask)
            rep2 = self.swe_pooling(chain2, chain2_mask)
            sim = F.cosine_similarity(rep1, rep2, dim=1).unsqueeze(-1)
            out = self.mlp(sim).squeeze(-1)
            
            return out
        
        else:  # "difference"
            rep1 = self.swe_pooling(chain1, chain1_mask)
            rep2 = self.swe_pooling(chain2, chain2_mask)
            diff = rep1 - rep2
            if self.normalize_difference:
                diff = F.normalize(diff, p=2, dim=1)
            out = self.mlp(diff).squeeze(-1)
            
            return out


class AffinityPredictionModel(pl.LightningModule):
    """
    Lightning wrapper for training. Siamese logic in training.py
    """
    def __init__(self, predictor: AffinityPredictor, learning_rate=1e-4):
        super().__init__()
        self.predictor = predictor
        self.learning_rate = learning_rate

        self.loss_fn = nn.MSELoss()
        self.pearson_corr = PearsonCorrCoef()
        self.spearman_corr = SpearmanCorrCoef()
        self.r2_score = R2Score()
        self.mse_metric = MeanSquaredError()

    def forward(self, chain1, chain1_mask, chain2, chain2_mask):
        return self.predictor(chain1, chain1_mask, chain2, chain2_mask)

    def training_step(self, batch, batch_idx):
        pass

    def validation_step(self, batch, batch_idx):
        pass

    def configure_optimizers(self):
        optimizer = AdamW(self.parameters(), lr=self.learning_rate)
        steps = self.trainer.estimated_stepping_batches
        scheduler = torch.optim.lr_scheduler.OneCycleLR(
            optimizer, max_lr=self.learning_rate, total_steps=steps
        )
        return {"optimizer": optimizer, "lr_scheduler": scheduler, "monitor": "val_loss"}


# Enhanced AffinityPredictor with Two-Head Architecture
# Add this to architectures.py

class DualHeadAffinityPredictor(nn.Module):
    """
    Enhanced AffinityPredictor with explicit two-head architecture.
    Simultaneously processes mutant and wildtype proteins to predict both ΔG and ΔΔG.
    
    embedding_method => "difference", "cosine", "cross_attention", or "cross_attention_swe".
    """
    def __init__(
        self,
        input_dim=1153,  # 1152 (ESM) + 1 (mutation)
        latent_dim=1024,
        num_slices=1024,
        num_ref_points=128,
        dropout_rate=0.2,
        freeze_swe=False,
        embedding_method="difference",
        normalize_difference=False,
        num_hidden_layers=2,
        # cross-attn
        num_cross_attn_layers=2,
        num_attention_heads=4,
        cross_ffn_dim=2048,
        use_dual_head=True,  # Enable dual-head by default
        ddg_signal_gain=1.0,  # Initial gain for ddG signal
        ddg_signal_multiplier=20.0, # FIXED: multiplier for ddG signal (vnew65.0)
    ):
        super().__init__()
        self.embedding_method = embedding_method.lower()
        self.normalize_difference = normalize_difference
        self.input_dim = input_dim
        self.use_dual_head = use_dual_head
        self._ddg_log_counter = 0
        
        # DEBUG: Confirm this version is running
        print(f"[MODEL INIT] DualHeadAffinityPredictor created: version={ARCH_VERSION}, dual_head={use_dual_head}, method={self.embedding_method}")
        
        # ESM dimension (without mutation channel)
        self.esm_dim = input_dim - 1  # 1152
        
        # Define cross-attention stack if needed
        self.cross_stack = None
        if "cross_attention" in self.embedding_method:
            self.cross_stack = MutationSpecificCrossAttentionStack(
                d_model=self.esm_dim,  # 1152
                num_heads=num_attention_heads,
                ffn_dim=cross_ffn_dim,
                dropout=dropout_rate,
                num_layers=num_cross_attn_layers
            )
        
        # Enhanced Mutation-Aware SWE Pooling
        self.swe_pooling = None
        if self.embedding_method in ["difference", "cosine", "cross_attention_swe"]:
            # For SWE, we use the full input_dim (1153)
            self.swe_pooling = MutationAwareSWEPooling(
                d_in=input_dim,
                num_slices=num_slices,
                num_ref_points=num_ref_points,
                freeze_swe=freeze_swe
            )
        
        # Define aggregator MLP in-dimensions
        if self.embedding_method == "cosine":
            in_features = 1
        else:
            in_features = num_slices  # difference or cross_attention_swe => [B, num_slices]
        
        # Add projection layer for cross_attention to avoid dynamic creation
        self.cross_attn_projection = None
        if self.embedding_method == "cross_attention":
            cross_proj_in = input_dim  # Full dimension including mutation channel
            cross_proj_out = in_features
            self.cross_attn_projection = nn.Linear(cross_proj_in, cross_proj_out, bias=False)
        
        # Define dG head (main prediction head)
        layers = []
        current_dim = in_features
        for _ in range(num_hidden_layers):
            layers.append(nn.Linear(current_dim, latent_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout_rate))
            current_dim = latent_dim
        layers.append(nn.Linear(current_dim, 1))
        self.dg_mlp = nn.Sequential(*layers)
        
        # Define ΔΔG head for direct prediction
        # FIX (vnew64.0): Shallow 2-layer MLP + residual skip connection.
        # DDGACT diagnostics from vnew62-63 showed 7-layer ReLU network causes 
        # progressive variance collapse (std: 0.1 → 7e-05). Each ReLU zeros ~50%
        # of activations, so 7 layers → 0.5^7 = 0.8% signal survival.
        # Solution: (1) 2 layers only, (2) skip connection preserves raw input signal.
        if self.use_dual_head:
            # Shallow nonlinear pathway (2 layers)
            self.ddg_hidden = nn.Sequential(
                nn.Linear(in_features, latent_dim),
                nn.GELU(),  # GELU instead of ReLU - no zero-capping, smoother gradients
                nn.Linear(latent_dim, latent_dim),
                nn.GELU(),
            )
            # Output projection
            self.ddg_out = nn.Linear(latent_dim, 1)
            # Skip connection: project input directly to output dimension
            self.ddg_skip = nn.Linear(in_features, 1)
        
        # Source type embedding for conditional inference
        # User-friendly types for inference:
        #   0 = "mutant"    - Single mutant predictions (most common)
        #   1 = "wt_pairs"  - Wildtype pairs with absolute binding affinity
        #   2 = "antibody"  - Antibody-antigen binding (CDR-focused)
        self.source_type_embedding = nn.Embedding(3, 32)  # 32-dim embedding
        self.source_type_projection = nn.Linear(in_features + 32, in_features)  # Project back to in_features

        # FIXED: Restoring the historical 'learnable gain' strategy
        # This allows the model to amplify the ddG signal early in Stage B
        self.ddg_signal_gain = nn.Parameter(torch.tensor(float(ddg_signal_gain)))
        self.ddg_signal_multiplier = float(ddg_signal_multiplier)

    def _extract_features(self, chain1, chain1_mask, chain2, chain2_mask):
        """
        Extract feature representation for a protein complex.
        Returns a vector representation suitable for prediction.
        """
        if "cross_attention" in self.embedding_method:
            # Process through mutation-specific cross-attention
            c1_out, c2_out = self.cross_stack(chain1, chain2, chain1_mask, chain2_mask)
            
            if self.embedding_method == "cross_attention_swe":
                # Apply enhanced SWE pooling
                rep1 = self.swe_pooling(c1_out, chain1_mask)  # [B, num_slices]
                rep2 = self.swe_pooling(c2_out, chain2_mask)  # [B, num_slices]
                
                # Difference aggregator
                diff = rep1 - rep2
                if self.normalize_difference:
                    diff = F.normalize(diff, p=2, dim=1)
                return diff
                
            elif self.embedding_method == "cross_attention":
                # Use mutation-weighted pooling
                # Extract mutation channel to guide pooling
                d_in = c1_out.shape[2]
                use_context = (d_in > 1153)
                if use_context:
                    c1_mut = c1_out[:, :, -2:-1]  # [B, L, 1]
                    c2_mut = c2_out[:, :, -2:-1]  # [B, L, 1]
                else:
                    c1_mut = c1_out[:, :, -1:]  # [B, L, 1]
                    c2_mut = c2_out[:, :, -1:]  # [B, L, 1]
                
                # Weighted pooling - gives higher weight to mutated positions
                c1_weights = F.softmax(c1_mut * 10, dim=1)  # Sharpen weights
                c2_weights = F.softmax(c2_mut * 10, dim=1)
                
                c1_pool = torch.sum(c1_out * c1_weights, dim=1)  # [B, 1153/1154]
                c2_pool = torch.sum(c2_out * c2_weights, dim=1)  # [B, 1153/1154]
                
                # Create difference representation
                diff = c1_pool - c2_pool
                if self.normalize_difference:
                    diff = F.normalize(diff, p=2, dim=1)
                
                # Use pre-defined projection layer instead of creating one dynamically
                if self.cross_attn_projection is not None:
                    diff = self.cross_attn_projection(diff)
                
                return diff
        
        elif self.embedding_method == "cosine":
            # Enhanced SWE => [B, num_slices]
            rep1 = self.swe_pooling(chain1, chain1_mask)
            rep2 = self.swe_pooling(chain2, chain2_mask)
            sim = F.cosine_similarity(rep1, rep2, dim=1).unsqueeze(-1)
            return sim
        
        else:  # "difference"
            rep1 = self.swe_pooling(chain1, chain1_mask)
            rep2 = self.swe_pooling(chain2, chain2_mask)
            diff = rep1 - rep2
            if self.normalize_difference:
                diff = F.normalize(diff, p=2, dim=1)
            return diff
    
    def _extract_residue_features(self, chain1, chain1_mask, chain2, chain2_mask):
        """
        Extract RESIDUE-LEVEL features (before pooling) for computing differences.
        Used for ddG to preserve mutation-specific information.
        
        Returns:
            c1_out, c2_out: [B, L1, D] and [B, L2, D] attended residue features
        """
        if "cross_attention" in self.embedding_method:
            c1_out, c2_out = self.cross_stack(chain1, chain2, chain1_mask, chain2_mask)
            return c1_out, c2_out
        else:
            # For non-cross-attention methods, return inputs directly
            return chain1, chain2
    
    def forward(self, mut_chain1, mut_chain1_mask, mut_chain2, mut_chain2_mask, 
            wt_chain1=None, wt_chain1_mask=None, wt_chain2=None, wt_chain2_mask=None,
            source_type_ids=None):
        """
        Dual-head forward method that can handle both modes:
        1. Standard mode: Just predict dG for mutant complex
        2. Dual-head mode: Predict both dG and direct ddG when wildtype is provided
        
        For ddG: Uses RESIDUE-LEVEL differences before pooling to preserve mutation info.
        ddG = ddg_mlp(pool(mut_features - wt_features)) instead of 
              ddg_mlp(pool(mut_features) - pool(wt_features))
        
        Args:
            source_type_ids: Optional[Tensor] of shape [B], values 0/1/2 for conditioning
        
        Returns:
            If wildtype inputs are None or use_dual_head=False:
                Returns mutant dG prediction only
            Else:
                Returns tuple of (mutant_dG, direct_ddG_prediction)
        """
        # ============== OPTIMIZED: Cache residue features ==============
        # Get mutant RESIDUE-LEVEL features first (used for both dG and ddG)
        mut_c1_res, mut_c2_res = self._extract_residue_features(
            mut_chain1, mut_chain1_mask, mut_chain2, mut_chain2_mask)
        
        # Pool for dG prediction (reuses cached residue features)
        if "cross_attention_swe" in self.embedding_method:
            rep1 = self.swe_pooling(mut_c1_res, mut_chain1_mask)
            rep2 = self.swe_pooling(mut_c2_res, mut_chain2_mask)
            mut_features = rep1 - rep2
            if self.normalize_difference:
                mut_features = F.normalize(mut_features, p=2, dim=1, eps=1e-8)
        else:
            # Fallback pooling
            mut_features = self._extract_features(mut_chain1, mut_chain1_mask, mut_chain2, mut_chain2_mask)
        
        # ============== SOURCE TYPE CONDITIONING ==============
        # Apply source type conditioning if provided
        if source_type_ids is not None:
            # Get source type embedding [B, 32]
            src_emb = self.source_type_embedding(source_type_ids)
            # Concatenate with features and project back
            conditioned_features = torch.cat([mut_features, src_emb], dim=-1)
            mut_features = self.source_type_projection(conditioned_features)
        
        # Predict dG for mutant
        dg_pred = self.dg_mlp(mut_features).squeeze(-1)
        
        # If no wildtype or dual head is disabled, just return mutant dG
        if not self.use_dual_head or wt_chain1 is None or wt_chain2 is None:
            return dg_pred
        
        # ============== RESIDUE-LEVEL ddG COMPUTATION ==============
        # Get wildtype RESIDUE-LEVEL features (mutant already cached above)
        wt_c1_res, wt_c2_res = self._extract_residue_features(
            wt_chain1, wt_chain1_mask, wt_chain2, wt_chain2_mask)
        
        
        # (Debug logging removed - was polluting training output)
        
        # Compute RESIDUE-LEVEL differences BEFORE pooling
        # This preserves mutation-specific changes at each position
        # Handle sequence length differences by taking minimum length
        L_c1 = min(mut_c1_res.shape[1], wt_c1_res.shape[1])
        L_c2 = min(mut_c2_res.shape[1], wt_c2_res.shape[1])
        
        c1_diff = mut_c1_res[:, :L_c1, :] - wt_c1_res[:, :L_c1, :]  # [B, L1, D]
        c2_diff = mut_c2_res[:, :L_c2, :] - wt_c2_res[:, :L_c2, :]  # [B, L2, D]
        
        # Update masks for truncated length
        c1_diff_mask = mut_chain1_mask[:, :L_c1] if mut_chain1_mask is not None else None
        c2_diff_mask = mut_chain2_mask[:, :L_c2] if mut_chain2_mask is not None else None
        
        # [DIFF CHECK] Diagnostic Logging
        # Measure cosine similarity at the approximate mutation site to check for embedding collapse.
        if self._ddg_log_counter % 200 == 1:
            with torch.no_grad():
                # Extract center position
                mid_idx = L_c1 // 2
                
                # Vectors at midpoint
                v_mut = mut_c1_res[:, mid_idx, :]
                v_wt = wt_c1_res[:, mid_idx, :]
                
                # Cosine similarity
                cos_sim = F.cosine_similarity(v_mut, v_wt, dim=1).mean().item()
                diff_norm = (v_mut - v_wt).norm(dim=1).mean().item()
                
                logger.info(f"[DIFF CHECK] Batch {self._ddg_log_counter}: CosSim at mid={cos_sim:.5f}, DiffNorm={diff_norm:.5f}")
                
        # NOW pool the differences
        if self.swe_pooling is not None:
            # =================================================================
            # HYBRID POOLING: Global SWE + Local Mutation-Site-Centric (vnew37.0)
            # This ensures local mutation signals are not diluted by global pooling
            # =================================================================
            
            # Concatenate chain differences along sequence dimension
            combined_diff = torch.cat([c1_diff, c2_diff], dim=1)  # [B, L_comb, D=1153]
            if c1_diff_mask is not None and c2_diff_mask is not None:
                combined_mask = torch.cat([c1_diff_mask, c2_diff_mask], dim=1)
            else:
                combined_mask = None

            # A. Global Component: Standard SWE pooling (capture global stability context)
            global_diff = self.swe_pooling(combined_diff, combined_mask) # [B, num_slices]

            # B. Local Component: Mutation-Site-Centric Pooling (MSCP)
            # CRITICAL FIX (v49.0): Extract indicator from RAW INPUT chains, NOT cross-attention output!
            # The cross-attention stack applies 0.9 convex combination at each layer (5 layers = 0.9^5 = 59% decay).
            # Using mut_chain1/mut_chain2 (raw inputs) instead of mut_c1_res/mut_c2_res (diluted outputs).
            
            # Determine if we have context window (1154-dim) or standard (1153-dim)
            d_raw = mut_chain1.shape[2]
            use_context = (d_raw > 1153)
            
            # Extract RAW indicator from INPUT chains (before cross-attention!)
            if use_context:
                c1_mut_raw = mut_chain1[:, :L_c1, -2:-1]
                c2_mut_raw = mut_chain2[:, :L_c2, -2:-1]
            else:
                c1_mut_raw = mut_chain1[:, :L_c1, -1:]
                c2_mut_raw = mut_chain2[:, :L_c2, -1:]
            
            mut_indicator = torch.cat([c1_mut_raw, c2_mut_raw], dim=1)  # [B, L_comb, 1]
            
            # Stability: Clamp indicator to be non-negative and bounded.
            # In case attention stack produced weird values, we force them back to [0, 2]
            mut_indicator = mut_indicator.clamp(min=0.0, max=2.0)
            
            # Weighted average focusing ONLY on the mutation sites
            # use 0.001 epsilon to avoid nan for WT samples (mut_sum=0)
            mut_sum = mut_indicator.sum(dim=1).clamp(min=1e-3)
            
            # MSCP Calculation with additional stability guard
            mscp_esm = (combined_diff[:, :, :1152] * mut_indicator).sum(dim=1) / mut_sum # [B, 1152]
            mscp_mut = (mut_indicator * mut_indicator).sum(dim=1) / mut_sum # [B, 1] (should be ~1.0)
            
            # Project local delta into the same slice space as global features 
            # using the SHARED theta projection and mut_projection from SWE
            # This ensures consistent representation between global and local paths
            local_diff_esm = self.swe_pooling.theta(mscp_esm)
            local_diff_mut = self.swe_pooling.mut_projection(mscp_mut)
            local_diff = local_diff_esm + local_diff_mut  # [B, num_slices]

            # C. Combine: Residual-style addition + Gain
            # FIXED: Apply signal_gain AFTER normalization so it actually has effect
            # Previously, L2-norm undid the scaling entirely.
            diff_multiplier = getattr(self, "ddg_signal_multiplier", 20.0)
            diff_features = (global_diff + local_diff) * diff_multiplier
            
            # ===================================================================
            # SIGNAL FLOW LOGGING: Track local vs global contribution
            if not hasattr(self, '_ddg_log_counter'):
                self._ddg_log_counter = 0
            self._ddg_log_counter += 1
            
            should_log = (self._ddg_log_counter % 200 == 1)
            if should_log:
                g_mag = global_diff.abs().mean().item()
                l_mag = local_diff.abs().mean().item()
                logger.info(f"[DDG SIGNAL] Batch {self._ddg_log_counter}: Global_mag={g_mag:.4f}, Local_mag={l_mag:.4f}")
                #region agent log
                try:
                    # Inspect what the model thinks the mutation indicator is (both tail channels)
                    d_raw_dbg = int(mut_chain1.shape[2])
                    # indicator candidate stats on chain1/chain2 for last and second-last channels
                    def _chan_stats(x):
                        return {
                            "min": float(x.min().item()),
                            "max": float(x.max().item()),
                            "mean": float(x.float().mean().item()),
                            "std": float(x.float().std().item()),
                        }
                    c1_last = _chan_stats(mut_chain1[:, :L_c1, -1])
                    c2_last = _chan_stats(mut_chain2[:, :L_c2, -1])
                    c1_last2 = _chan_stats(mut_chain1[:, :L_c1, -2]) if d_raw_dbg >= 1154 else None
                    c2_last2 = _chan_stats(mut_chain2[:, :L_c2, -2]) if d_raw_dbg >= 1154 else None
                    mut_sum_dbg = float(mut_sum.mean().item()) if "mut_sum" in locals() else None
                    payload = {
                        "sessionId": "debug-session",
                        "runId": "pre-fix",
                        "hypothesisId": "F",
                        "location": "architectures.py:DualHeadAffinityPredictor:mscp_indicator_debug",
                        "message": "Indicator channel stats (last vs second-last) to detect double-indicator / wrong channel selection",
                        "data": {
                            "ddg_log_counter": int(self._ddg_log_counter),
                            "d_raw": d_raw_dbg,
                            "use_context_flag": bool(use_context),
                            "c1_last": c1_last,
                            "c2_last": c2_last,
                            "c1_last2": c1_last2,
                            "c2_last2": c2_last2,
                            "mut_sum_mean": mut_sum_dbg,
                        },
                        "timestamp": int(time.time() * 1000),
                    }
                    with open("/Users/supantha/Documents/code_v2/protein/.cursor/debug.log", "a") as f:
                        f.write(json.dumps(payload, default=str) + "\n")
                    logger.info(f"[AGENTLOG MSCP] d_raw={d_raw_dbg} use_context={use_context} c1_last={c1_last} c1_last2={c1_last2} mut_sum_mean={mut_sum_dbg}")
                except Exception:
                    pass
                #endregion
            
            if self.normalize_difference:
                diff_features = F.normalize(diff_features, p=2, dim=1, eps=1e-8)
            
            # FIXED: Apply signal_gain AFTER normalization so it actually scales the output
            diff_features = diff_features * self.ddg_signal_gain
        else:
            # Fallback: mean pooling of differences
            if c1_diff_mask is not None:
                c1_diff = c1_diff * c1_diff_mask.unsqueeze(-1).float()
                c1_pool = c1_diff.sum(dim=1) / c1_diff_mask.sum(dim=1, keepdim=True).clamp(min=1)
            else:
                c1_pool = c1_diff.mean(dim=1)
            if c2_diff_mask is not None:
                c2_diff = c2_diff * c2_diff_mask.unsqueeze(-1).float()
                c2_pool = c2_diff.sum(dim=1) / c2_diff_mask.sum(dim=1, keepdim=True).clamp(min=1)
            else:
                c2_pool = c2_diff.mean(dim=1)
            diff_features = c1_pool - c2_pool
            if self.normalize_difference:
                diff_features = F.normalize(diff_features, p=2, dim=1)
        
        # DEBUG: Check for NaNs/Infs in diff_features
        if torch.isnan(diff_features).any() or torch.isinf(diff_features).any():
            print(f"[DEBUG MODEL] NaN/Inf in diff_features! Shape: {diff_features.shape}")
            if c1_diff_mask is not None:
                print(f"  c1_mask sum: {c1_diff_mask.sum(dim=1).min().item()}")
            print(f"  c1_diff nan: {torch.isnan(c1_diff).any().item()}")
            print(f"  c2_diff nan: {torch.isnan(c2_diff).any().item()}")

        # ============== SOURCE TYPE CONDITIONING FOR DDG ==============
        # Apply same source conditioning to diff_features for ddG prediction
        # This allows ddG head to learn different behaviors for different data sources
        if source_type_ids is not None:
            src_emb = self.source_type_embedding(source_type_ids)
            conditioned_diff = torch.cat([diff_features, src_emb], dim=-1)
            diff_features = self.source_type_projection(conditioned_diff)

        # MSCP Hybrid: Skip 20x gain since MSCP provides raw, un-diluted signal
        # vnew64.0: Shallow 2-layer GELU + skip connection to preserve variance
        ddg_hidden_out = self.ddg_hidden(diff_features)
        ddg_pred = (self.ddg_out(ddg_hidden_out) + self.ddg_skip(diff_features)).squeeze(-1)

        #region agent log
        # Diagnose "train variance but eval constant" which often indicates dropout-only variance or head collapse.
        try:
            if should_log:
                # Diff feature stats across the batch
                df = diff_features.detach()
                df_mean = float(df.mean().item()) if df.numel() else None
                df_std = float(df.std().item()) if df.numel() else None
                df_abs_mean = float(df.abs().mean().item()) if df.numel() else None
                # per-sample spread: average std over features (helps detect "all samples identical")
                df_per_sample_std = float(df.float().std(dim=1).mean().item()) if df.dim() == 2 and df.shape[0] > 0 else None

                # ddg_pred stats (across batch)
                p = ddg_pred.detach()
                p_mean = float(p.mean().item()) if p.numel() else None
                p_std = float(p.std().item()) if p.numel() else None

                # If dropout is active (training), a second forward pass should differ.
                p2_std = None
                p_diff_std = None
                if self.training:
                    p2 = (self.ddg_out(self.ddg_hidden(diff_features)) + self.ddg_skip(diff_features)).squeeze(-1).detach()
                    p2_std = float(p2.std().item()) if p2.numel() else None
                    p_diff_std = float((p2 - p).std().item()) if p2.numel() else None

                # Weight/bias norms (to detect collapse to near-zero weights or bias-only prediction)
                lin_layers = [m for m in self.ddg_hidden.modules() if isinstance(m, nn.Linear)] if hasattr(self, "ddg_hidden") else []
                w0 = lin_layers[0] if len(lin_layers) > 0 else None
                wL = lin_layers[-1] if len(lin_layers) > 0 else None
                w0_norm = float(w0.weight.detach().norm().item()) if w0 is not None else None
                wL_norm = float(wL.weight.detach().norm().item()) if wL is not None else None
                bL_norm = float(wL.bias.detach().norm().item()) if (wL is not None and wL.bias is not None) else None

                payload = {
                    "sessionId": "debug-session",
                    "runId": "pre-fix",
                    "hypothesisId": "I",
                    "location": "architectures.py:DualHeadAffinityPredictor:ddg_head_eval_vs_train",
                    "message": "ddG head collapse vs dropout-only variance diagnostics",
                    "data": {
                        "ddg_log_counter": int(self._ddg_log_counter),
                        "model_training": bool(self.training),
                        "normalize_difference": bool(getattr(self, "normalize_difference", False)),
                        "ddg_signal_gain": float(self.ddg_signal_gain.detach().item()) if hasattr(self, "ddg_signal_gain") else None,
                        "diff_features_mean": df_mean,
                        "diff_features_std": df_std,
                        "diff_features_abs_mean": df_abs_mean,
                        "diff_features_per_sample_std_mean": df_per_sample_std,
                        "ddg_pred_mean": p_mean,
                        "ddg_pred_std": p_std,
                        "ddg_pred2_std": p2_std,
                        "ddg_pred_repeat_diff_std": p_diff_std,
                        "ddg_w0_norm": w0_norm,
                        "ddg_wL_norm": wL_norm,
                        "ddg_bL_norm": bL_norm,
                    },
                    "timestamp": int(time.time() * 1000),
                }
                # Log first (before file write that may fail on cluster)
                logger.info(
                    f"[AGENTLOG DDGHEAD] train={self.training} df_std={df_std:.4f} df_ps_std={df_per_sample_std:.4f} "
                    f"pred_std={p_std:.4f} pred_mean={p_mean:.4f} rep_diff_std={p_diff_std if p_diff_std is not None else 'NA'} "
                    f"w0={w0_norm:.2f} wL={wL_norm:.2f} bL={bL_norm if bL_norm is not None else 'NA'}"
                )

                # Layerwise activation trace to locate where variance collapses (ReLU dead / dropout-only variance)
                layer_stats = []
                try:
                    with torch.no_grad():
                        x = diff_features.detach()
                        for li, layer in enumerate(self.ddg_hidden):
                            x = layer(x)
                            st = {
                                "i": int(li),
                                "t": layer.__class__.__name__,
                                "mean": float(x.mean().item()) if x.numel() else None,
                                "std": float(x.std().item()) if x.numel() else None,
                            }
                            if isinstance(layer, nn.ReLU):
                                st["zero_frac"] = float((x == 0).float().mean().item()) if x.numel() else None
                            layer_stats.append(st)
                    # Compact summary for logs (first 2 + last 2 layers)
                    compact = (layer_stats[:2] + (["..."] if len(layer_stats) > 4 else []) + layer_stats[-2:])
                    logger.info(f"[AGENTLOG DDGACT] train={self.training} layers={compact}")
                except Exception:
                    layer_stats = []
                # File write may fail on cluster - that's OK
                try:
                    payload["data"]["ddg_layer_stats"] = layer_stats
                    with open("/Users/supantha/Documents/code_v2/protein/.cursor/debug.log", "a") as f:
                        f.write(json.dumps(payload, default=str) + "\n")
                except Exception:
                    pass
        except Exception:
            pass
        #endregion

        # DEBUG: Check ddg_pred
        if torch.isnan(ddg_pred).any() or torch.isinf(ddg_pred).any():
             print(f"[DEBUG MODEL] NaN/Inf in ddg_pred!")
        
        return dg_pred, ddg_pred


class DualHeadAffinityPredictionModel(pl.LightningModule):
    """
    Lightning wrapper for dual-head training.
    """
    def __init__(self, predictor: DualHeadAffinityPredictor, learning_rate=1e-4, ddg_loss_weight=1.0):
        super().__init__()
        self.predictor = predictor
        self.learning_rate = learning_rate
        self.ddg_loss_weight = ddg_loss_weight

        self.loss_fn = nn.MSELoss()
        self.pearson_corr = PearsonCorrCoef()
        self.spearman_corr = SpearmanCorrCoef()
        self.r2_score = R2Score()
        self.mse_metric = MeanSquaredError()
        
        # Save hyperparameters for checkpointing
        self.save_hyperparameters(ignore=['predictor'])

    def forward(self, mut_chain1, mut_chain1_mask, mut_chain2, mut_chain2_mask, 
                wt_chain1=None, wt_chain1_mask=None, wt_chain2=None, wt_chain2_mask=None,
                source_type_ids=None):
        return self.predictor(mut_chain1, mut_chain1_mask, mut_chain2, mut_chain2_mask,
                             wt_chain1, wt_chain1_mask, wt_chain2, wt_chain2_mask,
                             source_type_ids=source_type_ids)

    def training_step(self, batch, batch_idx):
        # Mutant data
        (c1, m1, c2, m2, y_mut) = batch["mutant"]
        
        # Wildtype data with valid mask
        (cw1, w1m, cw2, w2m, y_wt) = batch["wildtype"]
        has_wt = batch["has_wt"]
        
        if self.predictor.use_dual_head and has_wt.sum() > 0:
            # For samples with wildtype available, use dual-head prediction
            valid_samples = has_wt.bool()
            
            # Get predictions for valid samples
            dg_pred, ddg_pred = self(
                c1[valid_samples], m1[valid_samples], 
                c2[valid_samples], m2[valid_samples],
                cw1[valid_samples], w1m[valid_samples], 
                cw2[valid_samples], w2m[valid_samples]
            )
            
            # Calculate losses for valid samples
            dg_loss = self.loss_fn(dg_pred, y_mut[valid_samples])
            
            # Calculate true ddG as difference between mutant and wildtype dG
            true_ddg = y_mut[valid_samples] - y_wt[valid_samples]
            ddg_loss = self.loss_fn(ddg_pred, true_ddg)
            
            # Combined loss with weighting
            loss = dg_loss + self.ddg_loss_weight * ddg_loss
            
            # Process remaining samples (without wildtype) with standard prediction
            if (~valid_samples).sum() > 0:
                standard_dg_pred = self(
                    c1[~valid_samples], m1[~valid_samples], 
                    c2[~valid_samples], m2[~valid_samples]
                )
                standard_loss = self.loss_fn(standard_dg_pred, y_mut[~valid_samples])
                
                # Add to total loss, weighted by proportion of samples
                n_valid = valid_samples.sum()
                n_total = len(valid_samples)
                loss = (n_valid / n_total) * loss + ((n_total - n_valid) / n_total) * standard_loss
        else:
            # Standard prediction for all samples
            dg_pred = self(c1, m1, c2, m2)
            loss = self.loss_fn(dg_pred, y_mut)
        
        # Log metrics
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        # Similar to training_step but with more comprehensive metrics
        (c1, m1, c2, m2, y_mut) = batch["mutant"]
        (cw1, w1m, cw2, w2m, y_wt) = batch["wildtype"]
        has_wt = batch["has_wt"]
        
        # Store predictions and targets for all samples
        all_dg_preds = []
        all_dg_targets = []
        all_ddg_preds = []
        all_ddg_targets = []
        
        if self.predictor.use_dual_head and has_wt.sum() > 0:
            valid_samples = has_wt.bool()
            
            # Dual-head prediction for samples with wildtype
            dg_pred, ddg_pred = self(
                c1[valid_samples], m1[valid_samples], 
                c2[valid_samples], m2[valid_samples],
                cw1[valid_samples], w1m[valid_samples], 
                cw2[valid_samples], w2m[valid_samples]
            )
            
            # Calculate true ddG
            true_ddg = y_mut[valid_samples] - y_wt[valid_samples]
            
            # Store predictions and targets
            all_dg_preds.append(dg_pred)
            all_dg_targets.append(y_mut[valid_samples])
            all_ddg_preds.append(ddg_pred)
            all_ddg_targets.append(true_ddg)
            
            # Process remaining samples with standard prediction
            if (~valid_samples).sum() > 0:
                standard_dg_pred = self(
                    c1[~valid_samples], m1[~valid_samples], 
                    c2[~valid_samples], m2[~valid_samples]
                )
                all_dg_preds.append(standard_dg_pred)
                all_dg_targets.append(y_mut[~valid_samples])
        else:
            # Standard prediction for all samples
            dg_pred = self(c1, m1, c2, m2)
            all_dg_preds.append(dg_pred)
            all_dg_targets.append(y_mut)
            
            # For samples with wildtype, calculate implicit ddG
            if has_wt.sum() > 0:
                valid_samples = has_wt.bool()
                wt_dg_pred = self(cw1[valid_samples], w1m[valid_samples], 
                                 cw2[valid_samples], w2m[valid_samples])
                
                implicit_ddg_pred = dg_pred[valid_samples] - wt_dg_pred
                true_ddg = y_mut[valid_samples] - y_wt[valid_samples]
                
                all_ddg_preds.append(implicit_ddg_pred)
                all_ddg_targets.append(true_ddg)
        
        # Concatenate all predictions and targets
        if all_dg_preds:
            all_dg_preds = torch.cat(all_dg_preds)
            all_dg_targets = torch.cat(all_dg_targets)
            
            # Calculate dG metrics
            dg_mse = self.mse_metric(all_dg_preds, all_dg_targets)
            dg_pearson = self.pearson_corr(all_dg_preds, all_dg_targets)
            dg_spearman = self.spearman_corr(all_dg_preds, all_dg_targets)
            dg_r2 = self.r2_score(all_dg_preds, all_dg_targets)
            
            # Log dG metrics
            self.log('val_dg_mse', dg_mse, on_epoch=True, prog_bar=True)
            self.log('val_dg_pearson', dg_pearson, on_epoch=True)
            self.log('val_dg_spearman', dg_spearman, on_epoch=True)
            self.log('val_dg_r2', dg_r2, on_epoch=True)
        
        # Calculate ddG metrics if available
        if all_ddg_preds:
            all_ddg_preds = torch.cat(all_ddg_preds)
            all_ddg_targets = torch.cat(all_ddg_targets)
            
            ddg_mse = self.mse_metric(all_ddg_preds, all_ddg_targets)
            ddg_pearson = self.pearson_corr(all_ddg_preds, all_ddg_targets)
            ddg_spearman = self.spearman_corr(all_ddg_preds, all_ddg_targets)
            ddg_r2 = self.r2_score(all_ddg_preds, all_ddg_targets)
            
            # Log ddG metrics
            self.log('val_ddg_mse', ddg_mse, on_epoch=True, prog_bar=True)
            self.log('val_ddg_pearson', ddg_pearson, on_epoch=True)
            self.log('val_ddg_spearman', ddg_spearman, on_epoch=True)
            self.log('val_ddg_r2', ddg_r2, on_epoch=True)
            
            # Combined validation metric for early stopping
            combined_metric = dg_mse + self.ddg_loss_weight * ddg_mse
            self.log('val_combined_metric', combined_metric, on_epoch=True)
        
        return {'val_dg_mse': dg_mse if 'dg_mse' in locals() else None,
                'val_ddg_mse': ddg_mse if 'ddg_mse' in locals() else None}

    def test_step(self, batch, batch_idx):
        # Similar to validation_step but returns more detailed metrics
        return self.validation_step(batch, batch_idx)

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate)
        scheduler = torch.optim.lr_scheduler.OneCycleLR(
            optimizer, max_lr=self.learning_rate, total_steps=self.trainer.estimated_stepping_batches
        )
        return {"optimizer": optimizer, "lr_scheduler": scheduler, "monitor": "val_combined_metric"}