import torch
# import talos
from torch import device, cuda, cat, stack
import torch.nn as nn
import utils.CNN_Layers as CustomLayers


class CNN_Net(nn.Module):   # , talos.utils.TorchHistory):
    def __init__(self, prps, final_layer_size=5):
        super(CNN_Net, self).__init__()
        self.final_layer_size = final_layer_size
        self.device = device('cuda:0' if cuda.is_available() else 'cpu')
        print("CNN Initialized. Using: " + str(self.device))

        # LAYERS
        print(f"CNN Model Initialization")
        self.conv1 = CustomLayers.Conv_elu_maxpool_drop(1, 192, (11, 13, 11), stride=(4,4,4), pool=True, prps=prps)
        self.conv2 = CustomLayers.Conv_elu_maxpool_drop(192, 384, (5, 6, 5), stride=(1,1,1), pool=True, prps=prps)
        self.conv3_mid_flow = CustomLayers.Mid_flow(384, 384, prps=prps)
        self.conv4_sepConv = CustomLayers.Conv_elu_maxpool_drop(384, 96,(3, 4, 3), stride=(1,1,1), pool=True, prps=prps,
                                                                sep_conv=True)
        self.conv5_sepConv = CustomLayers.Conv_elu_maxpool_drop(96, 48, (3, 4, 3), stride=(1, 1, 1), pool=True,
                                                                prps=prps, sep_conv=True)
        self.fc1 = CustomLayers.Fc_elu_drop(113568, 10, prps=prps, softmax=False)      # TODO, concatenate clinical data after this
        self.fc2 = CustomLayers.Fc_elu_drop(20, final_layer_size, prps=prps, softmax=True)  # For now this works as output layer, though may be incorrect
        self.fc_clinical1 = CustomLayers.Fc_elu_drop(2, 30, prps=prps, softmax=False)
        self.fc_clinical2 = CustomLayers.Fc_elu_drop(30,10, prps=prps, softmax=False)


    # FORWARDS
    def forward(self, x):
        clinical_data = x[1].to(torch.float32)
        x = x[0]
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3_mid_flow(x)
        x = self.conv4_sepConv(x)
        x = self.conv5_sepConv(x)

        # FLATTEN x
        flatten_size = x.size(1) * x.size(2) * x.size(3) * x.size(4)
        x = x.view(-1, flatten_size)

        x = self.fc1(x)

        # Clinical
        clinical_data = torch.transpose(clinical_data, 0, 1)
        clinical_data = self.fc_clinical1(clinical_data)
        clinical_data = self.fc_clinical2(clinical_data)

        x = cat((x, clinical_data), dim=1)

        x = self.fc2(x)
        return x