import torch
import torch.nn as nn
from torchvision import models

class L2CS(nn.Module):
    def __init__(self, backbone_name='resnet50', num_bins=90, pretrained=True):
        """
        L2CS-Net Teacher Model Wrapper.
        
        Args:
            backbone_name (str): ResNet variant to use ('resnet18', 'resnet34', 'resnet50').
            num_bins (int): Number of bins for pitch and yaw classification (default 90).
            pretrained (bool): Whether to use ImageNet weights for the backbone.
        """
        super(L2CS, self).__init__()
        
        # 1. Initialize Backbone
        if backbone_name == 'resnet18':
            self.backbone = models.resnet18(weights='IMAGENET1K_V1' if pretrained else None)
            feat_dim = 512
        elif backbone_name == 'resnet34':
            self.backbone = models.resnet34(weights='IMAGENET1K_V1' if pretrained else None)
            feat_dim = 512
        else: # Default resnet50
            self.backbone = models.resnet50(weights='IMAGENET1K_V2' if pretrained else None)
            feat_dim = 2048
            
        # Remove the original FC layer
        self.backbone = nn.Sequential(*(list(self.backbone.children())[:-1]))
        
        # 2. Dual-Branch Heads (Pitch & Yaw)
        self.fc_pitch = nn.Linear(feat_dim, num_bins)
        self.fc_yaw = nn.Linear(feat_dim, num_bins)
        
        # 3. Parameters for expectation calculation
        self.num_bins = num_bins
        # Bin range is typically -90 to 90 degrees or 0 to 180. 
        # L2CS usually uses -90 to 90 for a total of 180 degrees.
        # We'll create a register_buffer for the idx tensor (0 to num_bins-1)
        self.register_buffer('idx_tensor', torch.arange(num_bins).float())

    def forward(self, x):
        """
        Forward pass.
        Returns:
            pitch_logits, yaw_logits: Raw scores for each bin.
        """
        # Feature extraction
        features = self.backbone(x)
        features = features.view(features.size(0), -1)
        
        # Predict bins
        pitch_logits = self.fc_pitch(features)
        yaw_logits = self.fc_yaw(features)
        
        return pitch_logits, yaw_logits

    def get_angles(self, pitch_logits, yaw_logits):
        """
        Convert logits to continuous angles (degrees) using soft-argmax (expectation).
        Assumes bins represent -90 to 90 degrees with a 2-degree step for 90 bins.
        """
        pitch_softmax = nn.functional.softmax(pitch_logits, dim=1)
        yaw_softmax = nn.functional.softmax(yaw_logits, dim=1)
        
        # Expectation: sum(prob * idx)
        pitch_idx = torch.sum(pitch_softmax * self.idx_tensor, dim=1)
        yaw_idx = torch.sum(yaw_softmax * self.idx_tensor, dim=1)
        
        # Map bin index back to degrees: (idx * 2) - 90
        # (Assuming 90 bins covering 180 degrees)
        pitch_deg = pitch_idx * 2 - 90 
        yaw_deg = yaw_idx * 2 - 90
        
        return pitch_deg, yaw_deg

def load_teacher_model(path=None, backbone='resnet50', num_bins=90, device='cpu'):
    """
    Utility to load the teacher model with pre-trained weights.
    """
    model = L2CS(backbone_name=backbone, num_bins=num_bins, pretrained=True)
    if path:
        # Load custom weights if provided (from L2CS-Net repo)
        checkpoint = torch.load(path, map_location=device)
        # Handle cases where the state_dict is nested
        state_dict = checkpoint.get('state_dict', checkpoint)
        model.load_state_dict(state_dict, strict=False)
    
    model.to(device)
    model.eval()
    return model