--- Trains the neural network.
-- 
-- Uses data generated by @{data_generation_call}.
-- @module train

require 'optim'
local arguments = require 'Settings.arguments'
require 'Nn.masked_huber_loss'

local M = {}

--- Saves a neural net model to disk.
-- 
-- The model is saved to `arguments.model_path` and labelled with the epoch
-- number.
-- @param model the neural net to save
-- @param epoch the current epoch number
-- @param valid_loss the validation loss of the current network
-- @local
function  M:_save_model(model, epoch, valid_loss)

  local model_information = {}
  model_information.epoch = epoch
  model_information.valid_loss = valid_loss
  model_information.gpu = arguments.gpu

  local net_type_str = arguments.gpu and '_gpu' or '_cpu'
  local model_file_name = arguments.model_path .. '/epoch_' .. epoch .. net_type_str .. '.model'
  local information_file_name = arguments.model_path .. '/epoch_' .. epoch .. net_type_str .. '.info'
  torch.save(model_file_name, model)
  torch.save(information_file_name, model_information)
end

--- Function passed to torch's [optim package](https://github.com/torch/optim).
-- @param params_new the neural network params
-- @param inputs the neural network inputs
-- @param targets the neural network targets
-- @param mask the mask vectors used for the loss function
-- @return the masked Huber loss on `inputs` and `targets` 
-- @return the gradient of the loss function
-- @see masked_huber_loss
-- @local
local function feval(params_new, inputs, targets, mask)
  -- set x to x_new, if different
  -- (in this simple implementation, x_new will typically always point to x,
  -- so the copy is really useless)
  if M.params ~= params_new then
    M.params:copy(params_new)
  end

  M.grads:zero()
  local outputs = M.network:forward(inputs)  
  local loss = M.criterion:forward(outputs, targets, mask)

  -- backward
  local dloss_doutput = M.criterion:backward(outputs, targets)
  M.network:backward(inputs, dloss_doutput)

  return loss, M.grads
end

--- Trains the neural network.
-- @param network the neural network (see @{net_builder})
-- @param data_stream a @{data_stream|DataStream} object which provides the
-- training data
-- @param epoch_count the number of epochs (passes of the training data) to train for
function M:train(network, data_stream, epoch_count)

  M.network = network
  M.data_stream = data_stream

  M.params, M.grads = network:getParameters()
  M.criterion = MaskedHuberLoss()
  if(arguments.gpu) then
    M.criterion = M.criterion:cuda()
  end

  local state = {learningRate = arguments.learning_rate}
  local lossSum = 0
  local optim_func = optim.adam

  -- optimization loop
  local timer = torch.Timer()
  for epoch = 1, epoch_count do
    timer:reset()
    data_stream:start_epoch(epoch)
    lossSum = 0

    for i=1, data_stream:get_train_batch_count() do
      local inputs, targets, mask = data_stream:get_train_batch(i)
      assert(mask)
      local _, loss = optim_func(function (x) return feval(x, inputs, targets, mask) end, M.params, state)
      lossSum = lossSum + loss[1]
    end

    print(string.format("Training loss: %f", lossSum / data_stream.train_batch_count))

    --check validation loss
    local valid_loss_sum = 0
    for i=1, data_stream:get_valid_batch_count() do

      local inputs, targets, mask = data_stream:get_valid_batch(i)
      assert(mask)
      local outputs = M.network:forward(inputs)
      local loss = M.criterion:forward(outputs, targets, mask)
      valid_loss_sum = valid_loss_sum + loss 
    end

    local valid_loss = valid_loss_sum / data_stream.valid_batch_count
    print(string.format("Validation loss: %f", valid_loss))
    print('Epoch took: ', timer:time().real)

    --saving the model
    print(epoch)
    if epoch % arguments.save_epoch == 0 then
      print("SAVING MODEL")
      self:_save_model(network, epoch, valid_loss)
    end
  end
  --end of train loop

end

return M
