""" models from the segmentation_models_pytorch library
    install the library before using these models
"""
import torch
import torch.nn as nn
import segmentation_models_pytorch as smp


class SegModel(nn.Module):
    def __init__(self, architecture='Unet', 
                    encoder_name='mobilenet_v2', 
                    encoder_weights='imagenet',
                    encoder_depth=4,
                    in_channels=3, 
                    classes=4):
        super().__init__()

        """
        architecture = {
            Unet, UnetPlusPlus, EfficientUnetPlusPlus, DeepLabV3, DeepLabV3+
        }
        encoder_name = {
            resnet18 / 11M, resnet34 / 21M, resnet50 / 23M, resnet101 / 42M
            densenet121 / 6M, densenet169 / 12M, densenet201 / 18M
            efficientnet-b0 / 4M, -b1 / 6M, -b2 / 7M, -b3 / 10M, -b4 / 17M, -b5 / 28M
            mobilenet_v2 / 2M
        }
        encoder_weights = {
            "imagenet", None
        }
        in_channels = 3
        classes = 4
        """

        if encoder_depth == 4:
            decoder_channels = (256, 128, 64, 32)

        exec("""self.model = smp.%s(
            encoder_name='%s',
            encoder_weights='%s',
            encoder_depth=%s,
            decoder_channels=%s,
            in_channels=%s,
            classes=%s,
            activation=None,
        )
        """ % (architecture, encoder_name, encoder_weights, encoder_depth, decoder_channels, in_channels, classes))


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