------------------------------------------------------------------------
-- Adapted from github.com/jcjohnson/torch-rnn/pull/66/commits/5e30c1d54dc9ed1d152e4a55e9c438775f623ef7

--[[ ReverseSequence ]] --
-- Reverses a sequence on a given dimension.
-- Example: Given a tensor of torch.Tensor({{1,2,3,4,5}, {6,7,8,9,10})
-- nn.ReverseSequence(1):forward(tensor) would give: torch.Tensor({{6,7,8,9,10},{1,2,3,4,5}})
------------------------------------------------------------------------
local ReverseSequence, parent = torch.class("nn.ReverseSequence", "nn.Module")

function ReverseSequence:__init(dim,gpu)
    parent.__init(self)
    self.output = torch.Tensor()
    self.gradInput = torch.Tensor()
    self.outputIndices = torch.LongTensor()
    self.gradIndices = torch.LongTensor()
    self.typ = 'torch.CudaTensor'
    if gpu and (gpu < 1) then
      self.typ = 'torch.LongTensor'
    end
end

function ReverseSequence:reverseOutput(input)
  self.output:resizeAs(input)
  self.outputIndices:resize(input:size())
  local T = input:size(1)
  for x = 1, T do
      self.outputIndices:narrow(1, x, 1):fill(T - x + 1)
  end
  self.output:gather(input, 1, self.outputIndices:type(self.typ))
end

function ReverseSequence:updateOutput(input)
  input = input:transpose(1, 2)
  self:reverseOutput(input)
  self.output = self.output:transpose(1, 2)
  return self.output
end

function ReverseSequence:reverseGradOutput(gradOutput)
    self.gradInput:resizeAs(gradOutput)
    self.gradIndices:resize(gradOutput:size())
    local T = gradOutput:size(1)
    for x = 1, T do
        self.gradIndices:narrow(1, x, 1):fill(T - x + 1)
    end
    self.gradInput:gather(gradOutput, 1, self.gradIndices:type(self.typ))
end

function ReverseSequence:updateGradInput(inputTable, gradOutput)
  gradOutput = gradOutput:transpose(1, 2)
  self:reverseGradOutput(gradOutput)
  self.gradInput = self.gradInput:transpose(1, 2)
  return self.gradInput
end