require 'torch'
require 'nn'
require './util/LSTM'
require './util/ReverseSequence'
require('./util/OneHot.lua')


local Model, parent = torch.class('nn.Model', 'nn.Module')


local function create_lstm(self, reverse)
  local lstm = nn.Sequential()

  if reverse then lstm:add(nn.ReverseSequence(2,self.gpu)) end

  for i = 1, self.rnn_layers do
    local prev_dim = self.rnn_size
    if i == 1 then prev_dim = self.wordvec_dim end
    local rnn = nn.LSTM(prev_dim, self.rnn_size)
    rnn.remember_states = false
    table.insert(self.rnns, rnn)
    lstm:add(rnn)
    if self.dropout > 0 then
      lstm:add(nn.Dropout(self.dropout))
    end
  end

  if reverse then lstm:add(nn.ReverseSequence(2,self.gpu)) end

  return lstm
end


function Model:__init(opt)
  self.gpu = opt.gpu
  self.rnn_size = opt.rnn_size
  self.rnn_layers = opt.rnn_layers
  self.dropout = opt.dropout
  self.batchnorm = opt.batchnorm
  self.unidirectional = opt.unidirectional
  self.wordvec_dim  = #(opt.alphabet)  -- ACGT = 4
  self.cnn = opt.cnn
  self.rnn = opt.rnn
  self.cnn_filters = tablex.map(tonumber, opt.cnn_filters:split('-'))
  self.cnn_size = opt.cnn_size
  self.cnn_pool = opt.cnn_pool
  self.batch_size = opt.batch_size
  self.num_classes = opt.num_classes

  self.rnns = {}
  self.model = nn.Sequential()

  -- Create input embedding (we always use onehot-encoded matrix in this project)
  self.model:add(OneHot(self.wordvec_dim))

  ------------------------------------------
  -------- RNN and CNN-RNN Models ----------
  ------------------------------------------
  if self.rnn then
    local CNN = nn.Sequential()
    local RNN = nn.Sequential()

    -----------------------------
    ---------- CNN-RNN ----------
    -----------------------------
    if self.cnn then
      -- only use 1 layer of convolution in CNN-RNN
      local pad_size = math.floor(self.cnn_filters[1]/2)
      CNN:add(nn.SpatialZeroPadding(0,0,pad_size, pad_size))
      CNN:add(nn.TemporalConvolution(self.wordvec_dim, self.cnn_size, self.cnn_filters[1]))
      CNN:add(nn.ReLU())
      self.model:add(CNN)
      self.wordvec_dim = self.cnn_size
    end

    -----------------------------
    ------------ RNN ------------
    -----------------------------
    local fwd = create_lstm(self, false)
    fwd:add(nn.Mean(2)) -- take mean of output vectors over time dimension
    local bwd = create_lstm(self, true)
    bwd:add(nn.Mean(2)) -- take mean of output vectors over time dimension

    local concat = nn.ConcatTable()
    local output_size

    if self.unidirectional then
      concat:add(fwd) -- uese ConcatTable for consistency w/ b-lstm
      output_size = self.rnn_size
    else
      concat:add(fwd)
      concat:add(bwd)
      output_size = self.rnn_size*2
    end

    RNN:add(concat)
    RNN:add(nn.JoinTable(2))

    self.model:add(RNN)

    -- Create output classifier of (CNN-)RNN
    self.model:add(nn.Linear((output_size), self.num_classes))
    self.model:add(nn.LogSoftMax())

  ---------------------------------------
  -------------- CNN Model --------------
  ---------------------------------------
  else
    -- Need to compute output sequnece size to be fed to linear classifier
    local input_size = self.wordvec_dim

    -- Create layers
    for layer = 1,#self.cnn_filters-1 do
      self.model:add(nn.TemporalConvolution(input_size, self.cnn_size, self.cnn_filters[layer]))
      input_size = self.cnn_size
      self.model:add(nn.ReLU())
      self.model:add(nn.TemporalMaxPooling(self.cnn_pool,self.cnn_pool))
      if self.dropout > 0 then self.model:add(nn.Dropout(self.dropout)) end
    end

    -- Last layer of convolution
    self.model:add(nn.TemporalConvolution(input_size, self.cnn_size, self.cnn_filters[#self.cnn_filters]))
    self.model:add(nn.ReLU())
    if self.dropout > 0 then self.model:add(nn.Dropout(self.dropout)) end

    -- Max pool across entire sequence to get unfiform output size,
    -- and transpose (view) to feed into linear classifier
    self.model:add(nn.Max(2))
    self.model:add(nn.View(-1,self.cnn_size))

    -- Output classifier of CNN --
    self.model:add(nn.Linear(self.cnn_size,self.num_classes))
    self.model:add(nn.LogSoftMax())
  end

  print('-------- Model Architechture ----------')
  print(self.model)
end



-- Model Functions
function Model:updateOutput(input)
  return self.model:forward(input)
end

function Model:backward(input, gradOutput, scale)
  return self.model:backward(input, gradOutput, scale)
end

function Model:parameters()
  return self.model:parameters()
end

function Model:training()
  self.model:training()
  parent.training(self)
end

function Model:evaluate()
  self.model:evaluate()
  parent.evaluate(self)
end

function Model:resetStates()
  for i, rnn in ipairs(self.rnns) do rnn:resetStates() end
end

function Model:clearState()
  self.model:clearState()
end