"""ResNet model definition for phenology classification."""

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


class ResNetPhenologyClassifier(nn.Module):
    """ResNet50-based classifier for plant phenology phases."""
    
    def __init__(self, num_classes: int, pretrained: bool = True, freeze_backbone: bool = False):
        """
        Initialize the ResNet classifier.
        
        Args:
            num_classes: Number of phenology phase classes
            pretrained: Whether to use pretrained ImageNet weights
            freeze_backbone: Whether to freeze backbone layers
        """
        super(ResNetPhenologyClassifier, self).__init__()
        
        # Load pretrained ResNet50
        if pretrained:
            self.resnet = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
        else:
            self.resnet = models.resnet50(weights=None)
        
        # Freeze backbone if specified
        if freeze_backbone:
            for param in self.resnet.parameters():
                param.requires_grad = False
        
        # Replace final fully connected layer
        num_features = self.resnet.fc.in_features
        self.resnet.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(num_features, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes)
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass through the network.
        
        Args:
            x: Input tensor of shape (batch_size, 3, 224, 224)
            
        Returns:
            Output tensor of shape (batch_size, num_classes)
        """
        return self.resnet(x)
    
    def unfreeze_backbone(self):
        """Unfreeze all layers for fine-tuning."""
        for param in self.resnet.parameters():
            param.requires_grad = True
    
    def get_features(self, x: torch.Tensor) -> torch.Tensor:
        """
        Extract features before classification layer.
        
        Args:
            x: Input tensor
            
        Returns:
            Feature tensor
        """
        # Forward through all layers except fc
        x = self.resnet.conv1(x)
        x = self.resnet.bn1(x)
        x = self.resnet.relu(x)
        x = self.resnet.maxpool(x)
        
        x = self.resnet.layer1(x)
        x = self.resnet.layer2(x)
        x = self.resnet.layer3(x)
        x = self.resnet.layer4(x)
        
        x = self.resnet.avgpool(x)
        x = torch.flatten(x, 1)
        
        return x


def create_model(num_classes: int, pretrained: bool = True, device: Optional[str] = None) -> ResNetPhenologyClassifier:
    """
    Create and initialize a ResNet model.
    
    Args:
        num_classes: Number of output classes
        pretrained: Whether to use pretrained weights
        device: Device to place model on ('cuda' or 'cpu')
        
    Returns:
        Initialized model
    """
    if device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    model = ResNetPhenologyClassifier(num_classes, pretrained=pretrained)
    model = model.to(device)
    
    return model
