import os
import torch
import torch.nn as nn
from torchvision.models import resnet50, ResNet50_Weights

class GreeneryClassifier(nn.Module):
    """
    A ResNet-50 based classifier optimized for land-cover categorization.
    Wraps the torchvision ResNet-50 architecture and replaces the final 
    fully connected layer to match the EuroSAT class count.
    """
    def __init__(self, num_classes=10, pretrained=True):
        super(GreeneryClassifier, self).__init__()
        
        # Load pre-trained ResNet50
        weights = ResNet50_Weights.DEFAULT if pretrained else None
        self.model = resnet50(weights=weights)
        
        # Replace the final fully connected layer
        num_ftrs = self.model.fc.in_features
        self.model.fc = nn.Linear(num_ftrs, num_classes)

    def forward(self, x):
        """
        Forward pass for the model.
        Args:
            x (torch.Tensor): Input tensor of shape (Batch, 3, 64, 64).
        Returns:
            torch.Tensor: Logits for the 10 land-cover classes.
        """
        return self.model(x)

def load_model(weights_path=None, num_classes=10, device=None):
    """
    Utility function to initialize and optionally load weights into the classifier.
    Args:
        weights_path (str): Path to the .pth state dictionary.
        num_classes (int): Number of output classes.
        device (str, optional): Device to load the model onto ('cuda' or 'cpu').
                                Defaults to CUDA if available, else CPU.
    Returns:
        GreeneryClassifier: The initialized and loaded model.
    """
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    model = GreeneryClassifier(num_classes=num_classes)
    if weights_path and os.path.exists(weights_path):
        model.load_state_dict(torch.load(weights_path, map_location=device))
    model.to(device)
    return model