
import math
import warnings

import torch
import torch.nn as nn
import torch.autograd as autograd
import torch.nn.functional as F

from . import _cuda_

class SRUFunction (autograd.Function):

    @staticmethod
    def forward (ctx,*args):

        output,*variables = _cuda_.sru(*args)
        ctx.save_for_backward(*variables)

        return output

    @staticmethod
    def backward (ctx,d_output):

        return _cuda_.dsru(d_output,*ctx.saved_variables)

class SRU (nn.Module):

    def __init__ (self,
        hidden_size,
        dimensions=1,
        bias=True,
        batch_first=False,
        dropout=0,
        bidirectional=False):

        super(SRU,self).__init__()

        self.hidden_size = hidden_size
        self.dimensions = dimensions
        self.batch_first = batch_first
        self.dropout = nn.Dropout(p=dropout) if dropout else None
        self.bidirectional = bidirectional

        self.weights = nn.Parameter(torch.Tensor(dimensions * hidden_size,hidden_size))
        self.bias = nn.Parameter(torch.Tensor(hidden_size) if bias else torch.empty(0))

        self.reset_parameters()

    def reset_parameters (self):

        stdv = 1 / math.sqrt(self.hidden_size)

        for weight in self.parameters():
            weight.data.uniform_(-stdv,+stdv)

    def forward (self,input,hidden=None):

        if (not hidden is None) and (self.bidirectional): warnings.warn(
            "Providing a custom hidden state to a bidirectional RNN may"
            "produce invalid outputs and gradients.",
            RuntimeWarning)

        input = input.transpose(0,self.dimensions) if self.batch_first else input
        hidden = torch.empty(size=[0] * self.dimensions + input.shape[-2:],device=input.device) if hidden is None else hidden

        output = reversed((0,i-h) for i,h in zip(input.shape,hidden.shape))
        output = F.pad(hidden,[x for tuple in output for x in tuple])

        output = SRUFunction.apply(input,hidden,output,self.weights,self.bias,self.dimensions,self.bidirectional)
        output = output if self.dropout is None else self.dropout(output)

        return output
