local OneHot, parent = torch.class('OneHot', 'nn.Module')
-- adapted from https://github.com/karpathy/char-rnn/blob/master/util/OneHot.lua

function OneHot:__init(outputSize)
  parent.__init(self)
  self.outputSize = outputSize
  -- We'll construct one-hot encodings by using the index method to
  -- reshuffle the rows of an identity matrix. To avoid recreating
  -- it every iteration we'll cache it.
  self._eye = torch.zeros(outputSize+1,outputSize)
  self._eye[{{1,outputSize},{1,outputSize}}] = torch.eye(outputSize)
  self._eye[outputSize+1] = torch.zeros(outputSize)
end

function OneHot:updateOutput(input)
  self.output:resize(input:size(1), input:size(2), self.outputSize):zero()
  for i = 1,input:size(1) do
      self._eye = self._eye:float()
      local longInput = input[i]:long()
      self.output[i]:copy(self._eye:index(1, longInput))
  end
  return self.output
end