import torch
import torch.nn as nn
import torch.nn.functional as F

class WF_Net(nn.Module):
    """
    Implements Poppi et al.'s WF Model structure.
    Wraps a pre-trained ResNet-18 and dynamically applies 
    weight-space gating matrix multiplication during the forward step.
    """
    def __init__(self, original_model: nn.Module, num_classes: int):
        super().__init__()
        # Extract the sequence of blocks/layers L from the original model
        self.conv1 = original_model.conv1
        self.bn1 = original_model.bn1
        self.relu = original_model.relu
        self.maxpool = original_model.maxpool
        
        self.layer1 = original_model.layer1
        self.layer2 = original_model.layer2
        self.layer3 = original_model.layer3
        self.layer4 = original_model.layer4
        
        self.avgpool = original_model.avgpool
        self.fc = original_model.fc
        
        # Target layer for filtering: layer4 block 1 conv2
        # We extract its static tensor data out of the autograd parameter pool
        self.target_conv = self.layer4[1].conv2
        self.original_w = nn.Parameter(self.target_conv.weight.data.clone().detach(), requires_grad=False)
        
        # Require: Alpha gating matrix. Shape: (num_classes, out_channels)
        # Initialized to 1.5 as per Poppi et al.'s verbatim specification
        out_channels = self.original_w.shape[0]
        #self.alpha = nn.Parameter(torch.ones(num_classes, out_channels) * 1.5)
        self.alpha = nn.Parameter(torch.ones(num_classes, out_channels))

    def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor:
        # 1. Run through early sequence of layers undisturbed
        x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
                
        # Run layer4 block 0 normally
        x = self.layer4[0](x)
        
        # -------------------------------------------------------------
        # HERE IT IS: Save the structural skip connection (identity) 
        # BEFORE modifying features via block 1's convolutions
        # -------------------------------------------------------------
        identity = x 
                
        # Now enter layer4 block 1
        x = self.layer4[1].conv1(x)
        x = self.layer4[1].bn1(x)
        x = self.layer4[1].relu(x)
                
        # [Your Step 1 Masking Math happens right here...]
        batch_alpha = self.alpha[target_class_indices]
        mask = torch.sigmoid(batch_alpha).view(x.size(0), -1, 1, 1)
        
        # Run the functional convolution
        x = F.conv2d(
            x, 
            weight=self.original_w, 
            bias=self.target_conv.bias, 
            stride=self.target_conv.stride, 
            padding=self.target_conv.padding
        )
        
        # Apply your WF-Net channel mask
        x = x * mask 
        x = self.layer4[1].bn2(x)
                
        # -------------------------------------------------------------
        # HERE IT IS USED: Add the pristine identity back to the gated output
        # -------------------------------------------------------------
        x = self.layer4[1].relu(x + identity) 
                
        # Final Classification Head Sequence
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        y_out = self.fc(x)
                
        return y_out