from .swinv2_model import SwinTransformerV2
from .unet_swin_model import unet_swin
from .mim.mim import build_mim_model
from ..training.mim_utils import load_pretrained

import logging


def build_model(config,
                pretrain: bool = False,
                pretrain_method: str = 'mim',
                logger: logging.Logger = None):
    """
    Given a config object, builds a pytorch model.

    Returns:
        model: built model
    """

    if pretrain:

        if pretrain_method == 'mim':
            model = build_mim_model(config)
            return model

    encoder_architecture = config.MODEL.TYPE
    decoder_architecture = config.MODEL.DECODER

    if encoder_architecture == 'swinv2':

        logger.info(f'Hit encoder only build, building {encoder_architecture}')

        window_sizes = config.MODEL.SWINV2.PRETRAINED_WINDOW_SIZES

        model = SwinTransformerV2(
            img_size=config.DATA.IMG_SIZE,
            patch_size=config.MODEL.SWINV2.PATCH_SIZE,
            in_chans=config.MODEL.SWINV2.IN_CHANS,
            num_classes=config.MODEL.NUM_CLASSES,
            embed_dim=config.MODEL.SWINV2.EMBED_DIM,
            depths=config.MODEL.SWINV2.DEPTHS,
            num_heads=config.MODEL.SWINV2.NUM_HEADS,
            window_size=config.MODEL.SWINV2.WINDOW_SIZE,
            mlp_ratio=config.MODEL.SWINV2.MLP_RATIO,
            qkv_bias=config.MODEL.SWINV2.QKV_BIAS,
            drop_rate=config.MODEL.DROP_RATE,
            drop_path_rate=config.MODEL.DROP_PATH_RATE,
            ape=config.MODEL.SWINV2.APE,
            patch_norm=config.MODEL.SWINV2.PATCH_NORM,
            use_checkpoint=config.TRAIN.USE_CHECKPOINT,
            pretrained_window_sizes=window_sizes)

        if config.MODEL.PRETRAINED and (not config.MODEL.RESUME):
            load_pretrained(config, model, logger)

    else:

        errorMsg = f'Unknown encoder architecture {encoder_architecture}'

        logger.error(errorMsg)

        raise NotImplementedError(errorMsg)

    if decoder_architecture is not None:

        if encoder_architecture == 'swinv2':

            window_sizes = config.MODEL.SWINV2.PRETRAINED_WINDOW_SIZES

            model = SwinTransformerV2(
                img_size=config.DATA.IMG_SIZE,
                patch_size=config.MODEL.SWINV2.PATCH_SIZE,
                in_chans=config.MODEL.SWINV2.IN_CHANS,
                num_classes=config.MODEL.NUM_CLASSES,
                embed_dim=config.MODEL.SWINV2.EMBED_DIM,
                depths=config.MODEL.SWINV2.DEPTHS,
                num_heads=config.MODEL.SWINV2.NUM_HEADS,
                window_size=config.MODEL.SWINV2.WINDOW_SIZE,
                mlp_ratio=config.MODEL.SWINV2.MLP_RATIO,
                qkv_bias=config.MODEL.SWINV2.QKV_BIAS,
                drop_rate=config.MODEL.DROP_RATE,
                drop_path_rate=config.MODEL.DROP_PATH_RATE,
                ape=config.MODEL.SWINV2.APE,
                patch_norm=config.MODEL.SWINV2.PATCH_NORM,
                use_checkpoint=config.TRAIN.USE_CHECKPOINT,
                pretrained_window_sizes=window_sizes)

        else:

            raise NotImplementedError()

        if decoder_architecture == 'unet':

            num_classes = config.MODEL.NUM_CLASSES

            if config.MODEL.PRETRAINED and (not config.MODEL.RESUME):
                load_pretrained(config, model, logger)

            model = unet_swin(encoder=model, num_classes=num_classes)

        else:
            error_msg = f'Unknown decoder architecture: {decoder_architecture}'
            raise NotImplementedError(error_msg)

    return model