###### ------------------- základní model DP -----------------

### importy
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import torch.optim as optim
from torch.nn import init
from torch.utils.data import Dataset, DataLoader

from config import Config

class myConv(nn.Module):
    "Reference: https://github.com/PetraNovotna/MIL_based_PVC_localization/blob/master/net.py"
    def __init__(self, in_size, out_size, filter_size=3, stride=1, pad=None, do_batch=1, dov=0):
        super().__init__()

        pad = int((filter_size - 1) / 2)

        self.do_batch = do_batch
        self.dov = dov
        self.conv = nn.Conv1d(in_size, out_size, filter_size, stride, pad)
        self.bn = nn.BatchNorm1d(out_size, momentum=0.1) ## momentum 0.1 je default

        if self.dov > 0:
            self.do = nn.Dropout(dov)

    def forward(self, inputs):

        outputs = self.conv(inputs)
        if self.do_batch:
            outputs = self.bn(outputs)

        outputs = F.relu(outputs)

        if self.dov > 0:
            outputs = self.do(outputs)

        return outputs


class ResNet(nn.Module):

    """
    Reference: https://github.com/PetraNovotna/MIL_based_PVC_localization/blob/master/net.py
    """

    def __init__(self, levels=6, lvl1_size=6, input_size=3, output_size=2, convs_in_layer=2, init_conv=6,
                 filter_size=5, mil_solution='max'):
        super().__init__()
        self.mil_solution = mil_solution
        self.levels = levels ## počet residuálních bloků (= ResNet)
        self.lvl1_size = lvl1_size ## počet filtrů každé konvoluční vrstvy uvnitř 1. residuálního bloku - o násobky se zvětšuje hloubka v blocích
        self.input_size = input_size ## vstupní počet filtrů do self.init_conv = počet svodů EKG
        self.output_size = output_size ## !!!!!! není využito
        self.convs_in_layer = convs_in_layer ## počet konvolučních vrstev uvnitř jednotlivých bloků
        self.filter_size = filter_size ## velikost kernelu

        # self.t = 0.5 * np.ones(output_size)

        self.init_conv = myConv(input_size, init_conv, filter_size=filter_size) ## počet filtrů v úvodní konvoluční vrstvě - init_conv

        self.layers = nn.ModuleList()
        for lvl_num in range(self.levels):

            if lvl_num == 0:
                self.layers.append(myConv(init_conv, int(lvl1_size * (lvl_num + 1)), filter_size=filter_size, dov=Config.drop_out))
            else:
                self.layers.append(myConv(int(lvl1_size * (lvl_num)) + int(lvl1_size * (lvl_num)) + init_conv,
                                          int(lvl1_size * (lvl_num + 1)), filter_size=filter_size, dov=Config.drop_out))       ## konkatenace x0,x,y

            for conv_num_in_lvl in range(self.convs_in_layer - 1):
                self.layers.append(
                    myConv(int(lvl1_size * (lvl_num + 1)), int(lvl1_size * (lvl_num + 1)), filter_size=filter_size, dov=Config.drop_out))

        self.conv_final = nn.Conv1d(int(lvl1_size * (self.levels)) + int(lvl1_size * (self.levels)) + init_conv, 2, 1,
                                    1, 1)


        ## weigths initialization wih xavier method
        for i, m in enumerate(self.modules()):
            if isinstance(m, nn.Conv2d):
                init.xavier_normal_(m.weight)
                init.constant_(m.bias, 0)

    def forward(self, x, lens):

        x.requires_grad = True

        x = self.init_conv(x)

        x0 = x

        ## aply all convolutions
        layer_num = -1
        for lvl_num in range(self.levels):

            for conv_num_in_lvl in range(self.convs_in_layer):
                layer_num += 1
                if conv_num_in_lvl == 1:
                    y = x

                x = self.layers[layer_num](x)

            heatmap = x

            ## skip conection to previous layer and to the input
            x = torch.cat((F.avg_pool1d(x0, 2 ** lvl_num, 2 ** lvl_num), x, y), 1)

            x = F.max_pool1d(x, 2, 2)

        heatmap = F.max_pool1d(heatmap, 2, 2)  ## feature signal

        x = self.conv_final(x)

        if self.mil_solution == 'max':
            x = F.adaptive_max_pool1d(x, 1)
        elif self.mil_solution == 'mean':
            x = F.adaptive_avg_pool1d(x, 1)
        else:
            print('error mil solution')

        x = x.view(list(x.size())[:2])

        x = torch.softmax(x, dim=1)

        return x, heatmap


#############################################################################################
