import torch.nn as nn
from .backbones.ConvNeXt import create_convnext_model 
from .backbones.ConvNeXtV2 import create_convnext_v2_model
from .backbones.EfficientNet import crete_efficientnet_v2_model
from .backbones.DinoV2 import create_dinov2_model

class MelanomaClassifier(nn.Module):
    def __init__(self, model_name='convnext_tiny', num_classes=2, pretrained=True, in_22k=False, freeze_model=False):
        """
        Initialize the Melanoma Classification model
        
        Args:
            model_name: Name of the ConvNeXt model variant to use
            num_classes: Number of output classes (2 for binary melanoma classification)
            pretrained: Whether to use pretrained weights
        """
        super().__init__()
        
        if model_name.__contains__('convnext_'):
            self.model, self.num_features = create_convnext_model(model_name=model_name, pretrained=pretrained, in_22k=in_22k)
            self.model.head = nn.Linear(self.num_features, num_classes)  
        elif model_name.__contains__('efficientnet'):
            self.model, self.num_features  = crete_efficientnet_v2_model(model_name=model_name, num_classes=num_classes, pretrained=pretrained, in_22k=in_22k)
        elif model_name.__contains__('convnextv2'):
            self.model, self.num_features = create_convnext_v2_model(model_name=model_name, num_classes = num_classes, pretrained=pretrained, in_22k=in_22k)
        elif model_name.__contains__('dinov2'):
            self.model , self.num_features = create_dinov2_model(model_name=model_name, pretrained=pretrained, use_registers = True, freeze=freeze_model, register_buffer=None)
            self.model.head = nn.Linear(self.num_features, num_classes, bias=True)  
        else:
            raise ValueError(f"Unsupported model name: {model_name}")
     
    def forward(self, x):
        return self.model(x)