from abc import ABC
from pathlib import Path
from typing import Union

import torch
from torch import nn
import torch.nn.functional as F
import pytorch_lightning as pl
from lib.modules.utils import AutoPad, Interpolate

#
# Sub - Modules
###################


class ConvModule(nn.Module):

    @property
    def shape(self):
        x = torch.randn(self.in_shape).unsqueeze(0)
        output = self(x)
        return output.shape[1:]

    def __init__(self, in_shape, activation: nn.Module = nn.ELU, pooling_size=None, use_bias=True, use_norm=True,
                 dropout: Union[int, float] = 0,
                 conv_filters=64, conv_kernel=5, conv_stride=1, conv_padding=0):
        super(ConvModule, self).__init__()

        # Module Paramters
        self.in_shape = in_shape
        in_channels, height, width = in_shape[0], in_shape[1], in_shape[2]
        self.activation = activation()

        # Convolution Paramters
        self.padding = conv_padding
        self.stride = conv_stride

        # Modules
        self.dropout = nn.Dropout2d(dropout) if dropout else False
        self.pooling = nn.MaxPool2d(pooling_size) if pooling_size else False
        self.norm = nn.BatchNorm2d(in_channels, eps=1e-04, affine=False) if use_norm else False
        self.conv = nn.Conv2d(in_channels, conv_filters, conv_kernel, bias=use_bias,
                              padding=self.padding, stride=self.stride
                              )

    def forward(self, x):
        x = self.norm(x) if self.norm else x

        tensor = self.conv(x)
        tensor = self.dropout(tensor) if self.dropout else tensor
        tensor = self.pooling(tensor) if self.pooling else tensor
        tensor = self.activation(tensor)
        return tensor


class DeConvModule(nn.Module):

    @property
    def shape(self):
        x = torch.randn(self.in_shape).unsqueeze(0)
        output = self(x)
        return output.shape[1:]

    def __init__(self, in_shape, conv_filters=3, conv_kernel=5, conv_stride=1, conv_padding=0,
                 dropout: Union[int, float] = 0, autopad=False,
                 activation: Union[None, nn.Module] = nn.ReLU, interpolation_scale=None,
                 use_bias=True, normalize=False):
        super(DeConvModule, self).__init__()
        in_channels, height, width = in_shape[0], in_shape[1], in_shape[2]
        self.padding = conv_padding
        self.stride = conv_stride
        self.in_shape = in_shape
        self.conv_filters = conv_filters

        self.autopad = AutoPad() if autopad else False
        self.interpolation = Interpolate(scale_factor=interpolation_scale) if interpolation_scale else False
        self.norm = nn.BatchNorm2d(in_channels, eps=1e-04, affine=False) if normalize else False
        self.dropout = nn.Dropout2d(dropout) if dropout else False
        self.de_conv = nn.ConvTranspose2d(in_channels, self.conv_filters, conv_kernel, bias=use_bias,
                                          padding=self.padding, stride=self.stride)

        self.activation = activation() if activation else None

    def forward(self, x):
        x = self.norm(x) if self.norm else x
        x = self.dropout(x) if self.dropout else x
        x = self.autopad(x) if self.autopad else x
        x = self.interpolation(x) if self.interpolation else x

        tensor = self.de_conv(x)
        tensor = self.activation(tensor) if self.activation else tensor
        return tensor

    def size(self):
        return self.shape


class RecurrentModule(nn.Module):

    @property
    def shape(self):
        x = torch.randn(self.in_shape).unsqueeze(0)
        output = self(x)
        return output.shape[1:]

    def __init__(self, in_shape, hidden_size, num_layers=1, cell_type=nn.GRU, use_bias=True, dropout=0):
        super(RecurrentModule, self).__init__()
        self.use_bias = use_bias
        self.num_layers = num_layers
        self.in_shape = in_shape
        self.hidden_size = hidden_size
        self.dropout = dropout
        self.rnn = cell_type(self.in_shape[-1] * self.in_shape[-2], hidden_size,
                             num_layers=num_layers,
                             bias=self.use_bias,
                             batch_first=True,
                             dropout=self.dropout)

    def forward(self, x):
        tensor = self.rnn(x)
        return tensor