"""PoseConv3D PyTorch architecture for inference."""
import torch
import torch.nn as nn
import torch.nn.functional as F

class PoseConv3DEncoder(nn.Module):
    def __init__(self, in_channels: int = 1, embed_dim: int = 64):
        super().__init__()
        self.c1 = nn.Conv3d(in_channels, 32, kernel_size=(3, 5, 5), stride=(1, 2, 2), padding=(1, 2, 2))
        self.c2 = nn.Conv3d(32, 64, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
        self.c3 = nn.Conv3d(64, 128, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
        self.c4 = nn.Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1))
        self.fc = nn.Linear(256, embed_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """(B, C, T, H, W) -> (B, embed_dim) normalized embedding."""
        h = F.relu(self.c1(x))
        h = F.relu(self.c2(h))
        h = F.relu(self.c3(h))
        h = F.relu(self.c4(h))
        h = torch.cat([h.mean(dim=(2, 3, 4)), h.amax(dim=(2, 3, 4))], dim=1)
        z = self.fc(h)
        return z / torch.sqrt((z * z).sum(dim=1, keepdim=True) + 1e-12)