import math
import torch
from torch.autograd import Variable
from cupy.cuda import function
from pynvrtc.compiler import Program
from collections import namedtuple

###

kernel = '''
extern "C"
__global__ void recurrent_reverse_forget_mult(float *dst, const float *f, const float *x, int SEQ, int BATCH, int HIDDEN)
{
  /*
  Note: destination is assumed to be one timestep longer than f or x where dst[SEQ] = h_{-1}
  This means dst array has a separate index than that of f or x
  */
  int hid = blockIdx.x * blockDim.x + threadIdx.x;
  int bid = blockIdx.y * blockDim.y + threadIdx.y;
  if(hid >= HIDDEN || bid >= BATCH)
     return;
  //
  for (int ts = SEQ - 1; ts >= 0; ts--) {
     // Good sanity check for debugging - only perform additions to a zeroed chunk of memory
     // Addition seems atomic or near atomic - you should get incorrect answers if doubling up via threads
     // Note: the index i needs to be offset by one as f[0] (f_t) is used for dst[1] (h_t) etc

     // To move timesteps, we step HIDDEN * BATCH
     // To move batches, we move HIDDEN
     // To move neurons, we move +- 1
     // Note: dst[dst_i] = ts * 100 + bid * 10 + hid; is useful for debugging

     int i           = (ts - 0) * HIDDEN * BATCH + bid * HIDDEN + hid;
     int dst_i       = (ts - 0) * HIDDEN * BATCH + bid * HIDDEN + hid;
     int dst_iplus1  = (ts + 1) * HIDDEN * BATCH + bid * HIDDEN + hid;
     dst[dst_i]      = f[i] * x[i];
     dst[dst_i]      += (1 - f[i]) * dst[dst_iplus1];
  }
}

extern "C"
__global__ void bwd_recurrent_reverse_forget_mult(const float *h, const float *f, const float *x, const float *gh, float *gf, float *gx, float *ghinit, int SEQ, int BATCH, int HIDDEN)
{
  /*
  Note: h is assumed to be one timestep longer than f, x, gf, gx, or gh where dst[SEQ] = h_{-1}
  This means dst array has a separate index than that of f or x
  */
  int hid = blockIdx.x * blockDim.x + threadIdx.x;
  int bid = blockIdx.y * blockDim.y + threadIdx.y;
  if(hid >= HIDDEN || bid >= BATCH)
     return;
  //
  double running_f = 0;
  for (int ts = 1; ts < SEQ + 1; ts++) {
  //for (int ts = SEQ; ts > 0; ts--) {
     int i           = (ts - 1) * HIDDEN * BATCH + bid * HIDDEN + hid;
     int dst_i       = (ts - 0) * HIDDEN * BATCH + bid * HIDDEN + hid;
     int dst_iminus1 = (ts - 1) * HIDDEN * BATCH + bid * HIDDEN + hid;
     int dst_iplus1  = (ts + 1) * HIDDEN * BATCH + bid * HIDDEN + hid;
     //
     running_f       += gh[dst_iminus1];
     // Gradient of X
     gx[i]           = f[i] * running_f;
     // Gradient of F
     gf[i]           = (x[i] - h[dst_i]) * running_f;
     //
     // The line below is likely more numerically stable than (1 - f[i]) * running_f;
     running_f       = running_f - f[i] * running_f;
  }
  ghinit[bid * HIDDEN + hid] = running_f;
}
'''

###

class CPUReverseForgetMult(torch.nn.Module):
    def __init__(self):
        super(CPUReverseForgetMult, self).__init__()

    def forward(self, f, x, hidden_init=None):
        result = []
        ###
        # Reverse the forget and inputs so we traverse from end to beginning
        forgets = f.split(1, dim=0)[::-1]
        inputs = (f * x).split(1, dim=0)[::-1]
        prev_h = hidden_init
        for i, h in enumerate(inputs):
            h = h.squeeze()
            if prev_h is not None: h = h + (1 - forgets[i]) * prev_h
            result.append(h)
            prev_h = h
        ###
        # Reverse the results as it's currently [h_t, ... h_0]
        result = result[::-1]
        # Now it's [h_0, ..., h_t]
        ###
        return torch.cat(result, dim=0)


class GPUReverseForgetMult(torch.autograd.Function):
    configured_gpus = {}
    ptx = None
    def __init__(self):
        super(GPUReverseForgetMult, self).__init__()

    def compile(self):
        if self.ptx is None:
            program = Program(kernel.encode(), 'recurrent_forget_mult.cu'.encode())
            GPUReverseForgetMult.ptx = program.compile()

        if torch.cuda.current_device() not in GPUReverseForgetMult.configured_gpus:
            m = function.Module()
            m.load(bytes(self.ptx.encode()))

            self.forget_mult = m.get_function('recurrent_reverse_forget_mult')
            self.bwd_forget_mult = m.get_function('bwd_recurrent_reverse_forget_mult')

            Stream = namedtuple('Stream', ['ptr'])
            self.stream = Stream(ptr=torch.cuda.current_stream().cuda_stream)

            GPUReverseForgetMult.configured_gpus[torch.cuda.current_device()] = (self.forget_mult, self.bwd_forget_mult, self.stream)

        self.forget_mult, self.bwd_forget_mult, self.stream = GPUReverseForgetMult.configured_gpus[torch.cuda.current_device()]

    def forward(self, f, x, hidden_init=None):
        self.compile()
        seq_size, batch_size, hidden_size = f.size()
        result = f.new(seq_size + 1, batch_size, hidden_size).zero_()
        if hidden_init is not None: result[seq_size, :, :] = hidden_init
        ###
        grid = (math.ceil(hidden_size / 256), math.ceil(batch_size / 4))
        self.forget_mult(grid=grid, block=(256, 4), args=[result.data_ptr(), f.data_ptr(), x.data_ptr(), seq_size, batch_size, hidden_size], stream=self.stream)
        self.save_for_backward(f, x, hidden_init)
        self.result = result
        return result[:-1, :, :]

    def backward(self, grad_h):
        self.compile()
        f, x, hidden_init = self.saved_tensors
        h = self.result
        ###
        seq_size, batch_size, hidden_size = f.size()
        grad_f = f.new(*f.size()).zero_()
        grad_x = f.new(*f.size()).zero_()
        grad_h_init = f.new(batch_size, hidden_size).zero_()
        ###
        grid = (math.ceil(hidden_size / 256), math.ceil(batch_size / 4))
        self.bwd_forget_mult(grid=grid, block=(256, 4), args=[h.data_ptr(), f.data_ptr(), x.data_ptr(), grad_h.data_ptr(), grad_f.data_ptr(), grad_x.data_ptr(), grad_h_init.data_ptr(), seq_size, batch_size, hidden_size], stream=self.stream)
        ###
        if hidden_init is not None:
            return grad_f, grad_x, grad_h_init
        return grad_f, grad_x


class ReverseForgetMult(torch.nn.Module):
    r"""ForgetMult computes a simple recurrent equation:
    h_t = f_t * x_t + (1 - f_t) * h_{t-1}

    This equation is equivalent to dynamic weighted averaging.

    Inputs: X, hidden
        - X (seq_len, batch, input_size): tensor containing the features of the input sequence.
        - F (seq_len, batch, input_size): tensor containing the forget gate values, assumed in range [0, 1].
        - hidden_init (batch, input_size): tensor containing the initial hidden state for the recurrence (h_{t-1}).
        - use_cuda: If True, use the fast element-wise CUDA kernel for recurrence. If False, uses naive for loop. Default: True.
    """

    def __init__(self):
        super(ReverseForgetMult, self).__init__()

    def forward(self, f, x, hidden_init=None, use_cuda=True):
        # Use CUDA by default unless it's available
        use_cuda = use_cuda and torch.cuda.is_available()
        # Ensure the user is aware when ForgetMult is not GPU version as it's far faster
        if use_cuda: assert f.is_cuda and x.is_cuda, 'GPU ReverseForgetMult with fast element-wise CUDA kernel requested but tensors not on GPU'
        ###
        # Avoiding 'RuntimeError: expected a Variable argument, but got NoneType' when hidden_init is None
        if hidden_init is None: return GPUReverseForgetMult()(f, x) if use_cuda else CPUReverseForgetMult()(f, x)
        return GPUReverseForgetMult()(f, x, hidden_init) if use_cuda else CPUReverseForgetMult()(f, x, hidden_init)

###

if __name__ == '__main__':
    seq, batch, hidden = 35, 20, 650
    # Larger input (batch * seq * hidden) results in excessive memory for gradient check
    seq, batch, hidden = 10, 7, 19
    #seq, batch, hidden = 3, 7, 19
    #seq, batch, hidden = 1, 7, 19
    #seq, batch, hidden = 5, 1, 1
    a      = Variable(torch.rand(seq, batch, hidden).cuda(), requires_grad=True)
    forget = Variable(torch.rand(seq, batch, hidden).cuda(), requires_grad=True)
    last_h = Variable(torch.zeros(batch, hidden).cuda(), requires_grad=True)
    #last_h = Variable(torch.ones(batch, hidden).cuda() * 10, requires_grad=True)
    #a      = Variable(torch.ones(seq, batch, hidden).cuda(), requires_grad=True)
    #forget = Variable(torch.ones(seq, batch, hidden).cuda() * 0.5, requires_grad=True)

    print('CPU forget mult')
    print('=-=-' * 5)

    resulta = ReverseForgetMult()(forget, a, last_h, use_cuda=False)
    loss = resulta.pow(2).sum()
    loss.backward()
    #print(resulta.data)
    #print(forget.grad.data)

    #print(resulta.data)
    print('Result =', loss.data[0])
    print('X grad =', a.grad.mean().data[0])
    print('Forget grad =', forget.grad.mean().data[0])
    print('Last H grad =', last_h.grad.mean().data[0])

    x_grad_copy = a.grad.clone()

    print()
    print('CUDA forget mult')
    print('=-=-' * 5)

    a.grad.data *= 0
    forget.grad.data *= 0
    last_h.grad.data *= 0

    resultb = ReverseForgetMult()(forget, a, last_h, use_cuda=True)
    loss = resultb.pow(2).sum()
    loss.backward()
    #print(resultb.data)
    #print(forget.grad.data)

    print('Result =', loss.data[0])
    print('X grad =', a.grad.mean().data[0])
    print('Forget grad =', forget.grad.mean().data[0])
    print('Last H grad =', last_h.grad.mean().data[0])

    ###

    print()
    print('=-=-' * 5)
    print('(Xgrad - Xgrad).sum() =', (x_grad_copy - a.grad).sum().data[0])
    print('Residual error for result')
    print('=-=-' * 5)
    residual = (resulta - resultb)
    print(residual.abs().sum().data[0])
 
    # Had to loosen gradient checking, potentially due to general floating point badness?
    from torch.autograd import gradcheck
    inputs = [forget, a, last_h]
    test = gradcheck(ReverseForgetMult(), inputs, eps=1e-4, atol=1e-2)
    print(test)
