require 'torch'

require 'nn'

require 'optim'

require 'model'

include('util/auRoc.lua')

require 'lfs'

local cmd = torch.CmdLine()

-- model options

cmd:option('-init_from', '') -- resume previous model

cmd:option('-dropout', 0.5) -- dropout probability

cmd:option('-cnn', false) -- cnn model?

cmd:option('-rnn', false) -- rnn model?

cmd:option('-rnn_size', 32) -- rnn embedding size

cmd:option('-rnn_layers', 1) -- number of rnn layers

cmd:option('-unidirectional', false) -- unidirectional RNN

cmd:option('-cnn_filters', '9-5-3') -- convolution filter sizes at each layer for CNN. This parameter also specifies the number of CNN layers.

cmd:option('-cnn_pool', 2) -- covolution layer pool size

cmd:option('-cnn_size', 128) -- number of conv feature maps at each layer for CNN and CNN-RNN

-- Optimization options

cmd:option('-max_epochs', 50) -- number of iterations to run for

cmd:option('-learning_rate', 1e-2)

cmd:option('-grad_clip', 5) -- gradient clip value magnitude

cmd:option('-lr_decay_every', 10) -- learning rate decay iteration increment

cmd:option('-lr_decay_factor', 0.5) -- learning rate decay number

-- GPU

cmd:option('-gpu', 1) -- set to 0 if no GPU

-- Dataset options

cmd:option('-data_root', 'data') -- data root directory

cmd:option('-dataset', 'deepbind') -- dataset

cmd:option('-seq_length', 101) --length of DNA sequences

cmd:option('-TF', 'ATF1_K562_ATF1_-06-325-_Harvard') -- change for different TF

cmd:option('-alphabet', 'ACGT')

cmd:option('-size', 0) -- how much of each dataset to load. 0 = full

cmd:option('-batch_size', 256)

cmd:option('-class_labels','1,0') --specify positive label first

-- Other

cmd:option('-noprogressbar', false) -- lua progress bar

cmd:option('-name', '') --special name for model (optional)

cmd:option('-checkpoint_every', 0) -- save a model checkpoint every X iterations

-- Directory to save models to

cmd:option('-save_dir', 'models/')

local opt = cmd:parse(arg)

opt.class_labels_table = opt.class_labels:split(',')

opt.num_classes = #opt.class_labels_table

opt.alphabet_size = #opt.alphabet

local data_dir = opt.data_root..'/'..opt.dataset..'/'

-- Name of directory to save the models to

if opt.cnn and (not opt.rnn) then -- CNN

    model_name ='model=CNN,cnn_size='..opt.cnn_size..',cnn_filters='..opt.cnn_filters

elseif opt.cnn and opt.rnn then -- CNN-RNN

    model_name = 'model=CNN-RNN,cnn_size='..opt.cnn_size..',cnn_filter='..opt.cnn_filters:split('-')[1]..',rnn_size='..opt.rnn_size..',rnn_layers='..opt.rnn_layers

    if opt.unidirectional then model_name = model_name..',unidirectional' end

elseif (not opt.cnn) and opt.rnn then  -- RNN

    model_name = 'model=RNN,rnn_size='..opt.rnn_size..',rnn_layers='..opt.rnn_layers

    if opt.unidirectional then model_name = model_name..',unidirectional' end

else

  print('Need either -cnn or -rnn flag! Exiting')

  os.exit()

end

model_name = model_name..',dropout='..opt.dropout..',learning_rate='..opt.learning_rate..',batch_size='..opt.batch_size

if (opt.name ~= '') then model_name = model_name..','..opt.name end

-- Set up GPU stuff

local dtype = 'torch.FloatTensor'

if opt.gpu > 0  then

  collectgarbage()

  require 'cutorch'

  require 'cunn'

  cutorch.setDevice(opt.gpu )

  dtype = 'torch.CudaTensor'

  print(string.format('Running with CUDA on GPU %d', opt.gpu))

else

  print 'Running in CPU mode'

end

-- check if file exists

function file_exists(name)

   local f=io.open(name,"r")

   if f~=nil then io.close(f) return true else return false end

end

-- *Uncomment for loop and if statement to Loop through all TFs and run train/test*

-- for TF in lfs.dir(data_dir) do

-- if TF:sub(1,1) ~= '.' and TF:sub(1,3) ~= 'ALL' then

  flag = true

  opt.TF = TF or opt.TF

  opt.data_dir = data_dir..opt.TF

  print('------TF DATASET-------\n'..opt.data_dir..'\n')

  local save_file = opt.save_dir..model_name..'/'..opt.TF..'/log.t7'

  print('------MODEL LOCATION-------\n'..opt.save_dir..model_name..'\n')

  if file_exists(save_file) and opt.init_from == '' then

    print('Already trained! Exiting')

  else

    local log = {}

    log['train'] = {}

    log['test'] = {}

    log['best_epoch'] = 0

    log['best_auc'] = 0

    -- ============ Load Data =========== --

    require('data')

    data = {}

    data.train = createDatasetOneHot("train", opt)

    data.test = createDatasetOneHot("test", opt)

    train_size = data.train.inputs:size()

    test_size = data.test.inputs:size()

    -- ====== Initialize the model and criterion ======= --

    local model = nil

    if opt.init_from ~= '' then -- load a model

      print('Initializing from ', opt.init_from)

      model = torch.load(opt.init_from):type(dtype)

    else --create new model

      local opt_clone = torch.deserialize(torch.serialize(opt))

      print(opt_clone)

      model = nn.Model(opt_clone):type(dtype)

    end

    local params, grad_params = model:getParameters()

    local crit = nn.ClassNLLCriterion():type(dtype)

    local optim_config = {learningRate = opt.learning_rate, momentum = 0.9}

    local AUC = auRoc:new()

    -- ========== Run Train/Test ============ --

    local best_trainAUROC = 0

    for epoch = 1,opt.max_epochs do

      -- ========== Training ============ --

      print('======> Training epoch '..epoch)

      model:resetStates()

      model:training()

      for t = 1,train_size,opt.batch_size do

        if not opt.noprogressbar then

            xlua.progress(t, train_size)

        end

        model:resetStates()

        -- Loss function that we pass to an optim method

        local function f(w)

          assert(w == params)

          grad_params:zero()

          -- Get a minibatch

          x = data.train.inputs[t]:type(dtype)

          y = data.train.labels[t]:type(dtype)

          -- forward model

          local scores = model:forward(x)

          -- print(scores)

          -- os.exit()

          -- add scores and labels to AUC

          auc_in = scores[{{1,scores:size(1)},{1,1}}]:reshape(scores:size(1))

          for i = 1,auc_in:size(1) do AUC:add(math.exp(auc_in[i]), y[i]) end

          -- forward/backward criterion

          local loss = crit:forward(scores, y)

          local grad_scores = crit:backward(scores, y):view(x:size(1), 2, -1):reshape(x:size(1),2)

          -- backward model

          -- print(x)

          -- print(grad_scores)

          model:backward(x, grad_scores)

          -- clip gradients

          if opt.grad_clip > 0 then

            grad_params:clamp(-opt.grad_clip, opt.grad_clip)

          end

          return loss, grad_params

        end

        local _, loss = optim.adam(f, params, optim_config)

      end

      local trainAUROC = AUC:calculateAuc()

      print('\nTrain AUC: '..trainAUROC)

      AUC:zero()

      -- =========== Testing ============ --

      print('======> Testing epoch '..epoch)

      model:resetStates()

      model:evaluate()

      for t = 1,test_size,opt.batch_size do

        --progress bar

        if not opt.noprogressbar then

          xlua.progress(t, test_size)

        end

        -- get data

        x = data.test.inputs[t]:type(dtype)

        y = data.test.labels[t]:type(dtype)

        -- forward model

        model:resetStates()

        local scores = model:forward(x)

        -- add scores and labels to AUC

        auc_in = scores[{{1,scores:size(1)},{1,1}}]:reshape(scores:size(1))

        for i = 1,auc_in:size(1) do AUC:add(math.exp(auc_in[i]), y[i]) end

      end

      local testAUROC = AUC:calculateAuc()

      AUC:zero()

      print('\nTest AUC: '..testAUROC)

      -- ======== Checkpoint and Log ========== --

      -- check for training error

      if testAUROC < 0.1 then

        flag = false

        print('error in training, break from current TF')

        break

      end

      -- log the best AUC results and remember the model

      if epoch > 1 and (trainAUROC > best_trainAUROC) then

        best_trainAUROC = trainAUROC

        best_model=model:clone('weight','bias')

        log['best_epoch'] = epoch

        log['train_auc'] = trainAUROC

        log['test_auc'] = testAUROC

      end

      -- Save certain models

      if epoch % opt.checkpoint_every == 0 then

        torch.save(lfs.currentdir()..'/'..opt.save_dir..model_name..'/'..opt.TF..'/epoch'..epoch..'_model.t7', model)

      end

      -- decay learning rate

      if epoch % opt.lr_decay_every == 0 then

        local old_lr = optim_config.learningRate

        optim_config.learningRate = old_lr * opt.lr_decay_factor

      end

      -- Log every iteration (we don't log every model because of space)

      table.insert(log['test'],testAUROC)

      table.insert(log['train'],trainAUROC)

      collectgarbage()

    end

    -- ======== Save the Best model and Log ========== --

    if flag then

      lfs.mkdir(opt.save_dir..model_name)

      lfs.mkdir(opt.save_dir..model_name..'/'..opt.TF)

      torch.save(lfs.currentdir()..'/'..opt.save_dir..model_name..'/'..opt.TF..'/best_model.t7', model)

      torch.save(opt.save_dir..model_name..'/'..opt.TF..'/log'..'.t7', log)

    end

  end

-- end -- if directory

-- end -- loop through TFs