"""
PriviGaze Student Model - Ultra-Compact Gaze Estimation CNN

Design Philosophy:
- ~50-100K parameters total (target: fit on microcontrollers with TinyML)
- Inception blocks with factorized convolutions (exploits eye biology)
- Only takes light-corrected grayscale face as input (no eye crops)
- Designed for on-device inference with <1ms latency

Architecture inspired by:
- "One Eye is All You Need" (Athavale et al., 2022) - Inception for gaze
- DFT Gaze from GazeGen (281K params, distillation from 10x larger teacher)
- Eye biology: horizontal/vertical edge detectors mimic ocular muscle structure

Key design choices for disability support:
- Grayscale only: robust to varied lighting/occlusion
- Large receptive field: handles droopy eyes, head roll
- Factorized convolutions: 1x3 + 3x1 instead of 3x3 for efficiency
"""

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


def conv_bn_relu(in_ch, out_ch, kernel_size, stride=1, padding=0, groups=1):
    """Standard Conv-BN-ReLU block."""
    return nn.Sequential(
        nn.Conv2d(in_ch, out_ch, kernel_size, stride, padding, groups=groups, bias=False),
        nn.BatchNorm2d(out_ch),
        nn.ReLU(inplace=True),
    )


def depthwise_separable_conv(in_ch, out_ch, kernel_size, stride=1, padding=0):
    """Depthwise separable convolution for extreme parameter efficiency."""
    return nn.Sequential(
        # Depthwise
        nn.Conv2d(in_ch, in_ch, kernel_size, stride, padding, groups=in_ch, bias=False),
        nn.BatchNorm2d(in_ch),
        nn.ReLU(inplace=True),
        # Pointwise
        nn.Conv2d(in_ch, out_ch, 1, bias=False),
        nn.BatchNorm2d(out_ch),
        nn.ReLU(inplace=True),
    )


class FactorizedConvBlock(nn.Module):
    """Factorized 3x3 into 1x3 + 3x1 for efficiency.
    
    This mimics the ocular muscle structure: horizontal rectus (1x3) 
    and vertical rectus (3x1) for detecting gaze direction.
    """
    
    def __init__(self, in_ch, out_ch, stride=1):
        super().__init__()
        mid_ch = out_ch // 2
        
        self.horizontal = nn.Sequential(
            nn.Conv2d(in_ch, mid_ch, (1, 3), stride, (0, 1), bias=False),
            nn.BatchNorm2d(mid_ch),
            nn.ReLU(inplace=True),
        )
        
        self.vertical = nn.Sequential(
            nn.Conv2d(in_ch, out_ch - mid_ch, (3, 1), stride, (1, 0), bias=False),
            nn.BatchNorm2d(out_ch - mid_ch),
            nn.ReLU(inplace=True),
        )
    
    def forward(self, x):
        h = self.horizontal(x)
        v = self.vertical(x)
        return torch.cat([h, v], dim=1)


class InceptionBlock(nn.Module):
    """Lightweight inception block with factorized convolutions.
    
    Branches:
    1. 1x1 conv (pointwise - captures color/illumination)
    2. Factorized 1x3 + 3x1 (edge detection in H/V directions)
    3. MaxPool + 1x1 (spatial context)
    4. 2x 3x3 (standard feature extraction)
    
    All branches use depthwise separable convolutions for efficiency.
    """
    
    def __init__(self, in_ch, out_ch):
        super().__init__()
        # Each branch outputs out_ch // 4
        branch_ch = max(out_ch // 4, 8)
        
        # Branch 1: 1x1 pointwise
        self.branch1 = nn.Sequential(
            nn.Conv2d(in_ch, branch_ch, 1, bias=False),
            nn.BatchNorm2d(branch_ch),
            nn.ReLU(inplace=True),
        )
        
        # Branch 2: Factorized 1x3 + 3x1
        self.branch2 = FactorizedConvBlock(in_ch, branch_ch)
        
        # Branch 3: MaxPool + 1x1
        self.branch3 = nn.Sequential(
            nn.MaxPool2d(3, stride=1, padding=1),
            nn.Conv2d(in_ch, branch_ch, 1, bias=False),
            nn.BatchNorm2d(branch_ch),
            nn.ReLU(inplace=True),
        )
        
        # Branch 4: Two stacked 3x3 depthwise separable
        self.branch4 = nn.Sequential(
            depthwise_separable_conv(in_ch, branch_ch, 3, padding=1),
            depthwise_separable_conv(branch_ch, branch_ch, 3, padding=1),
        )
        
        # Final fusion
        total_ch = branch_ch * 4
        self.fusion = nn.Sequential(
            nn.Conv2d(total_ch, out_ch, 1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )
    
    def forward(self, x):
        b1 = self.branch1(x)
        b2 = self.branch2(x)
        b3 = self.branch3(x)
        b4 = self.branch4(x)
        out = torch.cat([b1, b2, b3, b4], dim=1)
        return self.fusion(out)


class LightCorrection(nn.Module):
    """Learnable light correction for grayscale face input.
    
    Applies per-channel (single channel for grayscale) affine transform
    to normalize lighting variations. This is critical for disability support
    where users may be in varied lighting conditions.
    """
    
    def __init__(self):
        super().__init__()
        # Learnable gamma correction parameter
        self.gamma = nn.Parameter(torch.ones(1))
        self.alpha = nn.Parameter(torch.ones(1))
        self.beta = nn.Parameter(torch.zeros(1))
    
    def forward(self, x):
        # x: [B, 1, H, W]
        # Apply gamma correction
        x = torch.pow(x.clamp(min=1e-6), self.gamma)
        # Apply affine transform
        x = self.alpha * x + self.beta
        return x


class PriviGazeStudent(nn.Module):
    """Ultra-compact student model for on-device gaze estimation.
    
    Input: 
        - face_gray: [B, 1, 224, 224] Light-corrected grayscale face
    
    Output:
        - pitch_pred: [B] gaze pitch in degrees
        - yaw_pred: [B] gaze yaw in degrees
        - features: [B, 128] feature representation (for distillation matching)
    
    Target: ~80K parameters, <1ms inference on mobile
    """
    
    def __init__(
        self,
        input_channels: int = 1,  # Grayscale
        feature_dim: int = 128,
        gaze_bins: int = 90,
    ):
        super().__init__()
        
        # Light correction
        self.light_correction = LightCorrection()
        
        # Stem: initial feature extraction
        self.stem = nn.Sequential(
            nn.Conv2d(input_channels, 32, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            # 32 x 112 x 112
            depthwise_separable_conv(32, 32, 3, stride=2, padding=1),
            # 32 x 56 x 56
        )
        
        # Stage 1: Inception blocks at high resolution
        self.stage1 = nn.Sequential(
            InceptionBlock(32, 64),
            # 64 x 56 x 56
            depthwise_separable_conv(64, 64, 3, stride=2, padding=1),
            # 64 x 28 x 28
        )
        
        # Stage 2: Deeper features
        self.stage2 = nn.Sequential(
            InceptionBlock(64, 96),
            # 96 x 28 x 28
            depthwise_separable_conv(96, 96, 3, stride=2, padding=1),
            # 96 x 14 x 14
        )
        
        # Stage 3: Abstract features
        self.stage3 = nn.Sequential(
            InceptionBlock(96, 128),
            # 128 x 14 x 14
            depthwise_separable_conv(128, 128, 3, stride=2, padding=1),
            # 128 x 7 x 7
        )
        
        # Stage 4: Global context
        self.stage4 = nn.Sequential(
            InceptionBlock(128, 160),
            # 160 x 7 x 7
            nn.AdaptiveAvgPool2d(1),
            # 160 x 1 x 1
        )
        
        # Feature projection
        self.feature_projection = nn.Sequential(
            nn.Flatten(),
            nn.Linear(160, feature_dim),
            nn.GELU(),
            nn.LayerNorm(feature_dim),
        )
        
        # Gaze regression heads (L2CS-Net style: per-angle binned regression)
        self.pitch_head = nn.Sequential(
            nn.Linear(feature_dim, feature_dim // 2),
            nn.GELU(),
            nn.Linear(feature_dim // 2, gaze_bins),
        )
        
        self.yaw_head = nn.Sequential(
            nn.Linear(feature_dim, feature_dim // 2),
            nn.GELU(),
            nn.Linear(feature_dim // 2, gaze_bins),
        )
        
        # Bin centers for expectation-based regression
        self.register_buffer(
            'bin_centers',
            torch.linspace(-90.0, 90.0, gaze_bins)
        )
        
        self.feature_dim = feature_dim
        self.gaze_bins = gaze_bins
    
    def forward(self, face_gray):
        """
        Args:
            face_gray: [B, 1, 224, 224] grayscale face image
        
        Returns:
            pitch_pred: [B]
            yaw_pred: [B]
            features: [B, feature_dim]
        """
        # Light correction
        x = self.light_correction(face_gray)
        
        # Stem
        x = self.stem(x)
        
        # Stages
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        
        # Feature projection
        features = self.feature_projection(x)
        
        # Gaze prediction
        pitch_logits = self.pitch_head(features)
        yaw_logits = self.yaw_head(features)
        
        pitch_probs = F.softmax(pitch_logits, dim=-1)
        yaw_probs = F.softmax(yaw_logits, dim=-1)
        
        pitch_pred = (pitch_probs * self.bin_centers).sum(dim=-1)
        yaw_pred = (yaw_probs * self.bin_centers).sum(dim=-1)
        
        return pitch_pred, yaw_pred, features
    
    def get_penultimate_features(self, face_gray):
        """Return features before regression heads for distillation."""
        _, _, features = self.forward(face_gray)
        return features


def count_parameters(model):
    """Count trainable parameters."""
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


if __name__ == "__main__":
    model = PriviGazeStudent()
    params = count_parameters(model)
    print(f"Student model parameters: {params:,}")
    
    # Test forward pass
    dummy = torch.randn(1, 1, 224, 224)
    pitch, yaw, features = model(dummy)
    print(f"Pitch: {pitch.shape}, Yaw: {yaw.shape}, Features: {features.shape}")