import argparse
import yaml

import torch
import torch.nn as nn

import timm

parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--cfg", default='configs/resnext50_32x4d.yaml', type=str)
args = parser.parse_args()

with open(args.cfg) as f:
    cfg = yaml.load(f, Loader=yaml.FullLoader)

class CustomResNext(nn.Module):
    def __init__(self, model_name='resnext50_32x4d', pretrained=False):
        super().__init__()
        self.model = timm.create_model(model_name, pretrained=pretrained)
        n_features = self.model.fc.in_features
        self.model.fc = nn.Linear(n_features, cfg['target_size'])

    def forward(self, x):
        x = self.model(x)
        return x
    
#### VIT
class VITModel(nn.Module):
    """
    Model Class for VIT Model
    """
    # Change num_classes
    def __init__(self, num_classes=cfg['target_size'], model_name='vit_base_patch16_384', pretrained=True):
        super(VITModel, self).__init__()
        self.model = timm.create_model(model_name, pretrained)
        self.model.head = nn.Linear(self.model.head.in_features, num_classes)

    def forward(self, x):
        x = self.model(x)
        return x
    
#### EfficientNet

class EfficientNetModel(nn.Module):
    """
    Model Class for EfficientNet Model
    """
    def __init__(self, num_classes=cfg['target_size'], model_name='efficientnet_b7', pretrained=True):
        super(EfficientNetModel, self).__init__()
        self.model = timm.create_model(model_name, pretrained=pretrained)
        self.model.classifier = nn.Linear(self.model.classifier.in_features, num_classes)
        
    def forward(self, x):
        x = self.model(x)
        return x

#### SeResNext   
class SeResNext(nn.Module):
    def __init__(self, num_classes=cfg['target_size'], model_name='seresnext26t_32x4d',  pretrained=True):
        super().__init__()
        self.model = timm.create_model(model_name, pretrained=pretrained)
        n_features = self.model.fc.in_features
        self.model.fc = nn.Linear(n_features, num_classes)
        

    def forward(self, x):
        x = self.model(x)
        return x