# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
import torch.nn.init as init

from torch import Tensor

class Transpose(nn.Module):
    """Wrapper class of torch.transpose() for Sequential module."""
    def __init__(self, shape: tuple):
        super(Transpose, self).__init__()
        
        self.shape = shape

    def forward(self, inputs: Tensor):
        return inputs.transpose(*self.shape)

class Linear(nn.Module):
    """
    Wrapper class of torch.nn.Linear
    Weight initialize by xavier initialization and bias initialize to zeros.
    """
    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
        super(Linear, self).__init__()
        
        self.linear = nn.Linear(in_features, out_features, bias=bias)
        init.xavier_uniform_(self.linear.weight)
        if bias:
            init.zeros_(self.linear.bias)

    def forward(self, x: Tensor) -> Tensor:
        return self.linear(x)


class View(nn.Module):
    """ Wrapper class of torch.view() for Sequential module. """
    def __init__(self, shape: tuple, contiguous: bool = False):
        super(View, self).__init__()
        
        self.shape = shape
        self.contiguous = contiguous

    def forward(self, inputs):
        if self.contiguous:
            inputs = inputs.contiguous()
        return inputs.view(*self.shape)

class ResidualConnectionModule(nn.Module):
    """
    Residual Connection module.
    output = (module(inputs) x module_factor + inputs x input_factor)
    """
    def __init__(
        self,
        module: nn.Module,
        module_factor: float = 1.0,
        input_factor: float = 1.0,
    ) -> None:
        super(ResidualConnectionModule, self).__init__()
        
        self.module = module
        self.module_factor = module_factor
        self.input_factor = input_factor

    def forward(self, inputs: Tensor) -> Tensor:
        return (self.module(inputs) * self.module_factor) + (inputs * self.input_factor)