import numpy as np

import torch

from torchvision import models

from efficientnet_pytorch import EfficientNet

import torch.nn.functional as F

import timm

def model_builder(architecture_name, output=6):

    if architecture_name.startswith("resnet"):

        net = eval("models." + architecture_name)(pretrained=True)

        net.fc = torch.nn.Linear(net.fc.in_features, output)

        return net

    elif architecture_name.startswith("efficientnet"):

        n = int(architecture_name[-1])

        net = EfficientNet.from_pretrained(f'efficientnet-b{n}')

        net._fc = torch.nn.Linear(net._fc.in_features, output)

        return net

    elif architecture_name.startswith("densenet"):

        net = eval("models." + architecture_name)(pretrained=True)

        net.classifier = torch.nn.Linear(net.classifier.in_features, output)

        return net

    elif architecture_name == "vgg19":

        net = models.vgg19_bn(pretrained=True)

        net.classifier[6] = torch.nn.Linear(net.classifier[6].in_features, output)

        return net

    elif architecture_name == "seresnext":

        net = timm.create_model('gluon_seresnext101_32x4d', pretrained=True)

        net.fc = torch.nn.Linear(net.fc.in_features, 6)

        return net

# https://github.com/pudae/kaggle-hpa/blob/master/losses/loss_factory.py

def binary_focal_loss(gamma=2, **_):

    def func(input, target):

        assert target.size() == input.size()

        max_val = (-input).clamp(min=0)

        loss = input - input * target + max_val + ((-max_val).exp() + (-input - max_val).exp()).log()

        invprobs = F.logsigmoid(-input * (target * 2 - 1))

        loss = (invprobs * gamma).exp() * loss

        return loss.mean()

    return func

class Windowing(torch.nn.Module):

    def __init__(self, u=1, epsilon=1e-3, window_length=50, window_width=130, transform="sigmoid"):

        """

        Practical Window Setting Optimization for Medical Image Deep Learning https://arxiv.org/pdf/1812.00572.pdf

        :param u: Upper bound for image values, e.g. 255

        :param epsilon:

        :param window_length:

        :param window_width:

        """

        super(Windowing, self).__init__()

        self.conv = torch.nn.Conv2d(in_channels=3, out_channels=3, kernel_size=1)

        self.u = u

        if transform == "sigmoid":

            weight = (2 / window_width) * np.log((u/epsilon) - 1)

            bias = (-2 * window_length / window_width) * np.log((u / epsilon) - 1)

            self.transform = self.sigmoid_transform

        else:  # Linear window

            weight = u / window_width

            bias = (-u / window_width) * (window_length - (window_width / 2))

            self.transform = self.linear_transform

        self.conv.weight = torch.nn.Parameter(weight * torch.ones_like(self.conv.weight))

        self.conv.bias = torch.nn.Parameter(bias * torch.ones_like(self.conv.bias))

    def linear_transform(self, x):

        return torch.relu(torch.max(x, torch.tensor(self.u)))

    def sigmoid_transform(self, x):

        return self.u * torch.sigmoid(x)

    def forward(self, img):

        return self.transform(self.conv(img))

class ResNetModel(torch.nn.Module):

    def __init__(self, step_train=False, output=6):

        super(ResNetModel, self).__init__()

        self.net = models.resnet50(pretrained=True)

        self.net.fc = torch.nn.Linear(self.net.fc.in_features, output)

        self.blocks = ["layer1", "layer2", "layer3", "layer4"]

        self.frozen_blocks = 4

        # Gradually unfreeze layers throughout training

        if step_train:

            for name, param in self.net.named_parameters():

                param.requires_grad_(False)

            self.unfreeze_layers()

    def phase1_model(self):

        self.net.fc = torch.nn.Linear(self.net.fc.in_features, 1)

    def phase2_model(self):

        self.net.fc = torch.nn.Linear(self.net.fc.in_features, 5)

    def unfreeze_layers(self, lower_bound=0):

        for name, param in self.net.named_parameters():

            if self.frozen_blocks < 0:

                param.requires_grad_(True)

            elif name.split(".")[0] in ["fc"]:

                param.requires_grad_(True)

            elif name.split(".")[0] in self.blocks[self.frozen_blocks:]:

                param.requires_grad_(True)

        if self.frozen_blocks >= lower_bound:

            self.frozen_blocks -= 1

    def forward(self, x):

        return self.net(x)

class DenseNetModel(torch.nn.Module):

    def __init__(self, step_train=False, output=6):

        super(DenseNetModel, self).__init__()

        self.net = models.densenet169(pretrained=True)

        self.net.classifier = torch.nn.Linear(self.net.classifier.in_features, output)

        self.blocks = ["denseblock1", "denseblock2", "denseblock3", "denseblock4"]

        self.frozen_blocks = 4

        # Gradually unfreeze layers throughout training

        if step_train:

            for name, param in self.net.named_parameters():

                param.requires_grad_(False)

            self.unfreeze_layers()

    def phase1_model(self):

        self.net.fc = torch.nn.Linear(self.net.fc.in_features, 1)

    def phase2_model(self):

        self.net.fc = torch.nn.Linear(self.net.fc.in_features, 5)

    def unfreeze_layers(self, lower_bound=0):

        for name, param in self.net.named_parameters():

            if self.frozen_blocks < 0:

                param.requires_grad_(True)

            elif name.split(".")[0] in ["fc"]:

                param.requires_grad_(True)

            elif name.split(".")[0] in self.blocks[self.frozen_blocks:]:

                param.requires_grad_(True)

        if self.frozen_blocks >= lower_bound:

            self.frozen_blocks -= 1

    def forward(self, x):

        return self.net(x)

class EfficientNetModel(torch.nn.Module):

    """

    # Coefficients:   width,depth,res,dropout

    'efficientnet-b0': (1.0, 1.0, 224, 0.2),

    'efficientnet-b1': (1.0, 1.1, 240, 0.2),

    'efficientnet-b2': (1.1, 1.2, 260, 0.3),

    'efficientnet-b3': (1.2, 1.4, 300, 0.3),

    'efficientnet-b4': (1.4, 1.8, 380, 0.4),

    'efficientnet-b5': (1.6, 2.2, 456, 0.4),

    'efficientnet-b6': (1.8, 2.6, 528, 0.5),

    'efficientnet-b7': (2.0, 3.1, 600, 0.5),

    """

    def __init__(self, n=0, step_train=False, output=6):

        super(EfficientNetModel, self).__init__()

        self.net = EfficientNet.from_pretrained(f'efficientnet-b{n}')

        self.net._fc = torch.nn.Linear(self.net._fc.in_features, output)

        filters = [block._block_args.output_filters for block in self.net._blocks]

        self.freeze_points = (np.where(np.diff(filters) > 0)[0])  # 6 main block groups which can be frozen/unfrozen

        self.frozen_blocks = 6

        # Gradually unfreeze layers throughout training

        if step_train:

            for name, param in self.net.named_parameters():

                param.requires_grad_(False)

            self.unfreeze_layers()

    def phase1_model(self):

        self.net._fc = torch.nn.Linear(self.net._fc.in_features, 1)

    def phase2_model(self):

        self.net._fc = torch.nn.Linear(self.net._fc.in_features, 5)

    def unfreeze_layers(self, lower_bound=3):

        try:

            fp = self.freeze_points[self.frozen_blocks]

        except IndexError:

            fp = np.Inf

        for name, param in self.net.named_parameters():

            if name.split(".")[0] in ["_conv_head", "_bn1", "_fc"]:

                param.requires_grad_(True)

            elif name.split(".")[1].isnumeric():

                block_number = int(name.split(".")[1])

                if block_number > fp:

                    param.requires_grad_(True)

        if self.frozen_blocks >= lower_bound:

            self.frozen_blocks -= 1

            print("Trainable blocks:", 6 - self.frozen_blocks)

    def forward(self, x):

        return self.net(x)