import timm
import torch
import torch.nn as nn
 
EFFICIENTNET_DIM = 1280
 
 
class TemporalAttention(nn.Module):
    """
    Input  : (B, N_FRAMES, EFFICIENTNET_DIM)
    Output : (pooled (B, EFFICIENTNET_DIM), attn_weights (B, N_FRAMES))
    """
 
    def __init__(self, input_dim: int = EFFICIENTNET_DIM):
        super().__init__()
        self.attention = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.Tanh(),
            nn.Linear(128, 1),
        )
 
    def forward(self, x: torch.Tensor):
        scores  = self.attention(x)               # (B, 16, 1)
        weights = torch.softmax(scores, dim=1)     # (B, 16, 1)
        pooled  = (weights * x).sum(dim=1)         # (B, 1280)
        
        return pooled, weights.squeeze(-1)         # (B, 1280), (B, 16)
        
        
 
 
class FallDetector(nn.Module):
    """
    EfficientNet-Lite0 + Temporal Attention + MLP classifier.
 
    Forward:
        (B, 16, 3, 224, 224)
        -> per-frame EfficientNet-Lite0 -> (B, 16, 1280)
        -> TemporalAttention            -> (B, 1280), (B, 16)
        -> MLP classifier               -> (B, 1) logit
    Sigmoid is applied explicitly at inference, not inside the module.
    """
 
    def __init__(self, pretrained_backbone: bool = False):
        super().__init__()
        backbone = timm.create_model("efficientnet_lite0", pretrained=pretrained_backbone)
 
        self.conv_stem   = backbone.conv_stem
        self.bn1         = backbone.bn1
        self.blocks      = backbone.blocks
        self.conv_head   = backbone.conv_head
        self.bn2         = backbone.bn2
        self.global_pool = backbone.global_pool
 
        self.temporal_attention = TemporalAttention(EFFICIENTNET_DIM)
 
        self.classifier = nn.Sequential(
            nn.Linear(EFFICIENTNET_DIM, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(512, 128),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(128, 1),
        )
 
    def forward(self, x: torch.Tensor):
        B, T, C, H, W = x.shape
        x_flat = x.view(B * T, C, H, W)
 
        f = self.conv_stem(x_flat)
        f = self.bn1(f)
        f = self.blocks(f)
        f = self.conv_head(f)
        f = self.bn2(f)
        f = self.global_pool(f)                    # (B*16, 1280)
        f = f.view(B, T, EFFICIENTNET_DIM)          # (B, 16, 1280)
 
        pooled, attn_weights = self.temporal_attention(f)
        logit = self.classifier(pooled)             # (B, 1)
        
        return logit, attn_weights

    
    def count_parameters(self):
        total     = sum(p.numel() for p in self.parameters())
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        
        return total, trainable