import torch
import torch.nn as nn

class DualPoolMiniConv8x8(nn.Module):
    def __init__(self):
        super(DualPoolMiniConv8x8, self).__init__()
        # Using padding=1 for 8x8 to preserve spatial information at the boundaries
        self.conv = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1), nn.ReLU(inplace=True),
            nn.BatchNorm2d(16),
            nn.Conv2d(16, 32, kernel_size=3, padding=1), nn.ReLU(inplace=True),
            nn.BatchNorm2d(32),
            nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(inplace=True),
            nn.BatchNorm2d(64)
        )
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

    def forward(self, x):
        # x: (Batch, 1, 8, 8)
        x = self.conv(x)
        # Output: (Batch, 128)
        return torch.cat([self.avg_pool(x), self.max_pool(x)], dim=1).flatten(1)

class LIPEV2Ablation8x8(nn.Module):
    def __init__(self):
        super(LIPEV2Ablation8x8, self).__init__()
        self.app_net = DualPoolMiniConv8x8()
        
        # Geometry network (same as 16x16)
        self.geo_net = nn.Sequential(
            nn.Linear(956, 256),
            nn.LayerNorm(256),
            nn.ReLU(inplace=True),
            nn.Linear(256, 256),
            nn.ReLU(inplace=True)
        )
        
        # 64*2*4 (appearance) + 256 (geometry) = 512 + 256 = 768
        self.post_concat_bn = nn.BatchNorm1d(768)
        
        self.fusion = nn.Sequential(
            nn.Linear(768, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.1),
            nn.Linear(256, 128),
            nn.ReLU(inplace=True)
        )
        
        self.pitch_head = nn.Linear(128, 90)
        self.yaw_head = nn.Linear(128, 90)

    def forward(self, patches, landmarks):
        # patches: (Batch, 4, 8, 8)
        # landmarks: (Batch, 956)
        batch_size = patches.shape[0]
        
        # Process 4 patches per sample
        app_feat = self.app_net(patches.view(-1, 1, 8, 8)).view(batch_size, -1)
        geo_feat = self.geo_net(landmarks)
        
        combined = torch.cat([app_feat, geo_feat], dim=1)
        combined = self.post_concat_bn(combined)
        fused = self.fusion(combined)
        
        return self.pitch_head(fused), self.yaw_head(fused)