"""
Model utilities for telecom site classification
Handles ConvNeXt model loading and adaptation for transfer learning
"""

import torch
import torch.nn as nn
import timm
import os
from typing import Dict, Any, Optional, Tuple

class TelecomClassifier(nn.Module):
    """
    ConvNeXt-based telecom site classifier
    Uses transfer learning from food detection model
    """
    def __init__(self, num_classes: int = 3, pretrained: bool = True):
        super(TelecomClassifier, self).__init__()
        self.backbone = timm.create_model(
            'convnext_large.fb_in22k_ft_in1k', 
            pretrained=pretrained, 
            num_classes=0  # Remove classification head
        )
        self.feature_dim = self.backbone.num_features
        self.classifier = nn.Sequential(
            nn.LayerNorm(self.feature_dim),
            nn.Linear(self.feature_dim, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(512, 128),
            nn.ReLU(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(128, num_classes)
        )
        self._init_classifier_weights()
    def _init_classifier_weights(self):
        for module in self.classifier.modules():
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                nn.init.constant_(module.bias, 0)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        features = self.backbone(x)
        output = self.classifier(features)
        return output
    def freeze_backbone(self):
        for param in self.backbone.parameters():
            param.requires_grad = False
        print("🔒 Backbone frozen for transfer learning")
    def unfreeze_backbone(self):
        for param in self.backbone.parameters():
            param.requires_grad = True
        print("🔓 Backbone unfrozen for fine-tuning")
    def get_parameter_count(self) -> Dict[str, int]:
        backbone_params = sum(p.numel() for p in self.backbone.parameters())
        classifier_params = sum(p.numel() for p in self.classifier.parameters())
        total_params = backbone_params + classifier_params
        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
        return {
            'backbone': backbone_params,
            'classifier': classifier_params,
            'total': total_params,
            'trainable': trainable_params
        }

def load_model(
    model_path: str,
    num_classes: int = 3,
    device: str = 'cpu'
) -> Tuple[TelecomClassifier, Dict[str, Any]]:
    """
    Load trained telecom classifier model
    Args:
        model_path: Path to saved model
        num_classes: Number of output classes
        device: Device to load model on
    Returns:
        Tuple of (model, model_info)
    """
    print(f"📂 Loading model from {model_path}")
    model = TelecomClassifier(num_classes=num_classes, pretrained=False)
    checkpoint = torch.load(model_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    model_info = checkpoint.get('model_info', {})
    model_info['best_acc'] = checkpoint.get('best_acc', 'Unknown')
    model_info['epoch'] = checkpoint.get('epoch', 'Unknown')
    print(f"✅ Model loaded successfully")
    print(f"   Best accuracy: {model_info.get('best_acc', 'Unknown')}")
    print(f"   Epoch: {model_info.get('epoch', 'Unknown')}")
    return model, model_info