import os
import torch
import torch.nn as nn
import torchvision.models as models
from modelscope.msdatasets import MsDataset
from utils import url_download, create_dir, DOMAIN


def get_backbone(ver, backbone_list):
    for bb in backbone_list:
        if ver == bb['ver']:
            return bb

    print('Backbone name not found, using default option - alexnet.')
    return backbone_list[0]


def model_info(m_ver):
    backbone_list = MsDataset.load(
        'monetjoe/cv_backbones',
        subset_name='ImageNet1k_v1',
        split='train'
    )
    backbone = get_backbone(m_ver, backbone_list)
    m_type = str(backbone['type'])
    input_size = int(backbone['input_size'])
    return m_type, input_size


def download_model(log_name='vit_b_16_mel_2024-01-07_05-16-24'):
    pre_model_url = f'{DOMAIN}{log_name}/save.pt'
    pre_model_path = f"./model/{log_name}.pt"
    m_ver = '_'.join(log_name.split('_')[:-3])
    create_dir('./model')

    if not os.path.exists(pre_model_path):
        url_download(pre_model_url, pre_model_path)

    return pre_model_path, m_ver


def Classifier(cls_num: int, output_size: int, linear_output: bool):
    q = (1.0 * output_size / cls_num) ** 0.25
    l1 = int(q * cls_num)
    l2 = int(q * l1)
    l3 = int(q * l2)

    if linear_output:
        return torch.nn.Sequential(
            nn.Dropout(),
            nn.Linear(output_size, l3),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(l3, l2),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(l2, l1),
            nn.ReLU(inplace=True),
            nn.Linear(l1, cls_num)
        )

    else:
        return torch.nn.Sequential(
            nn.Dropout(),
            nn.Conv2d(output_size, l3, kernel_size=(1, 1), stride=(1, 1)),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d(output_size=(1, 1)),
            nn.Flatten(),
            nn.Linear(l3, l2),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(l2, l1),
            nn.ReLU(inplace=True),
            nn.Linear(l1, cls_num)
        )


class EvalNet():
    model = None
    m_type = 'squeezenet'
    input_size = 224
    output_size = 512

    def __init__(self, log_name, cls_num=4):
        saved_model_path, m_ver = download_model(log_name)
        self.m_type, self.input_size = model_info(m_ver)

        if not hasattr(models, m_ver):
            print('Unsupported model.')
            exit()

        self.model = eval('models.%s()' % m_ver)
        linear_output = self._set_outsize()
        self._set_classifier(cls_num, linear_output)
        checkpoint = torch.load(saved_model_path, map_location='cpu')
        if torch.cuda.is_available():
            checkpoint = torch.load(saved_model_path)

        self.model.load_state_dict(checkpoint, False)
        self.model.eval()

    def _set_outsize(self, debug_mode=False):
        for name, module in self.model.named_modules():
            if str(name).__contains__('classifier') or str(name).__eq__('fc') or str(name).__contains__('head'):
                if isinstance(module, torch.nn.Linear):
                    self.output_size = module.in_features
                    if debug_mode:
                        print(
                            f"{name}(Linear): {self.output_size} -> {module.out_features}")
                    return True

                if isinstance(module, torch.nn.Conv2d):
                    self.output_size = module.in_channels
                    if debug_mode:
                        print(
                            f"{name}(Conv2d): {self.output_size} -> {module.out_channels}")
                    return False

        return False

    def _set_classifier(self, cls_num, linear_output):
        if hasattr(self.model, 'classifier'):
            self.model.classifier = Classifier(
                cls_num, self.output_size, linear_output)
            return

        elif hasattr(self.model, 'fc'):
            self.model.fc = Classifier(
                cls_num, self.output_size, linear_output)
            return

        elif hasattr(self.model, 'head'):
            self.model.head = Classifier(
                cls_num, self.output_size, linear_output)
            return

        self.model.heads.head = Classifier(
            cls_num, self.output_size, linear_output)

    def forward(self, x):
        if torch.cuda.is_available():
            x = x.cuda()
            self.model = self.model.cuda()

        if self.m_type == 'googlenet' and self.training:
            return self.model(x)[0]
        else:
            return self.model(x)