require 'nn'

-- Returns a list of tensors for each line in the file

local function loadData(name, maxLoad)

    print(name)

    if maxLoad == 0 then maxLoad = 1000000000 end

    local data = seq.lines(name):take(maxLoad):copy()

    function data:size() return #data end

    return data

end

function table_invert(t)

   local s={}

   for k,v in pairs(t) do s[v]=k end

   return s

end

function createDatasetOneHot(typ, opt)

    if (not opt.size) then opt.size = 0 end

    if (not opt.batch_size) then opt.batch_size = 1 end

    local seqs = loadData( path.join(opt.data_dir, typ..'.fa'), opt.size)

    local size

    if opt.size == 0 then size = seqs:size() end

    local inputs = {}

    local outputs = {}

    local alphabet = opt.alphabet

    local rev_lookup = {}

    for i = 1,#alphabet do

      rev_lookup[alphabet:sub(i,i)] = i

    end

    local class_labels = table_invert(opt.class_labels_table)

    setmetatable(inputs, {__index = function(self, ind)

      local upper_limit = math.min((ind+opt.batch_size),((size/2)+1))

      local batch_len = upper_limit-ind

      local seq_len = opt.seq_length

      -- len = #seqs[ind*2]

      matrix = torch.zeros(batch_len, seq_len):fill(#alphabet+1)

      local batch_dim = 1

      for b = ind,upper_limit-1 do

        local str=seqs[b*2] -- have to multiply by 2 because of the way the data is set up with >1 and then sequence

        for i = 1,math.min(#str,seq_len),1 do

          if rev_lookup[str:sub(i,i)] then

            matrix[batch_dim][i] = rev_lookup[str:sub(i,i)]

          end

        end

        batch_dim = batch_dim+1

      end

      return matrix

    end})

    setmetatable(outputs, {__index = function(self, ind)

      local upper_limit = math.min((ind+opt.batch_size),((size/2)+1))

      local batch_len = upper_limit-ind

      local labels = torch.ones(batch_len)

      local batch_dim = 1

      for i = ind,upper_limit-1 do

        local line = seqs[(i*2)-1]:gsub(' ','')--get label from i*2 -1 (because of FASTA format)

        local label = line:split('>')[1]

        label = class_labels[label]

        labels[batch_dim] = torch.Tensor({label})

        batch_dim = batch_dim+1

      end

      return labels

    end})

    function inputs:size() return size/2 end

    function outputs:size() return size/2 end

    return {inputs=inputs, labels=outputs}

end