-- Francois Fleuret's Torch Toolbox

require 'torch'
require 'nn'

----------------------------------------------------------------------

colors = sys.COLORS

function printf(f, ...)
   print(string.format(f, unpack({...})))
end

function printfc(c, f, ...)
   printf(c .. string.format(f, unpack({...})) .. colors.black)
end

function logCommand(c)
   print(colors.blue .. '[' .. c .. '] -> [' .. sys.execute(c) .. ']' .. colors.black)
end

----------------------------------------------------------------------
-- Environment variables

defaultNbThreads = 1
defaultUseGPU = false

if os.getenv('TORCH_NB_THREADS') then
   defaultNbThreads = os.getenv('TORCH_NB_THREADS')
   print('Environment variable TORCH_NB_THREADS is set and equal to ' .. defaultNbThreads)
else
   print('Environment variable TORCH_NB_THREADS is not set, default is ' .. defaultNbThreads)
end

if os.getenv('TORCH_USE_GPU') then
   defaultUseGPU = os.getenv('TORCH_USE_GPU') == 'yes'
   print('Environment variable TORCH_USE_GPU is set and evaluated as ' .. tostring(defaultUseGPU))
else
   print('Environment variable TORCH_USE_GPU is not set, default is ' .. tostring(defaultUseGPU))
end

----------------------------------------------------------------------

function fftbInit(cmd, params)

   torch.setnumthreads(params.nbThreads)
   torch.setdefaulttensortype('torch.FloatTensor')
   torch.manualSeed(params.seed)

   -- Logging

   if params.rundir == '' then
      params.rundir = cmd:string('experiment', params, { })
   end

   paths.mkdir(params.rundir)

   if not params.noLog then
      -- Append to the log if there is one
      cmd:log(io.open(params.rundir .. '/log', 'a'), params)
   end

   -- Dealing with the CPU/GPU

   ffnn = {}

   -- By default, ffnn returns the entries from nn
   local mt = {}
   function mt.__index(table, key)
      return (cudnn and cudnn[key]) or (cunn and cunn[key]) or nn[key]
   end
   setmetatable(ffnn, mt)

   -- These are the tensors that can be kept on the CPU
   ffnn.SlowTensor = torch.Tensor
   ffnn.SlowStorage = torch.Storage
   -- These are the tensors that should be moved to the GPU
   ffnn.FastTensor = torch.Tensor
   ffnn.FastStorage = torch.Storage

   if params.useGPU then
      require 'cutorch'
      require 'cunn'
      require 'cudnn'

      if params.fastGPU then
         cudnn.benchmark = true
         cudnn.fastest = true
      end

      ffnn.FastTensor = torch.CudaTensor
      ffnn.FastStorage = torch.CudaStorage
   end
end

----------------------------------------------------------------------

function dimAtThatPoint(model, input)
   if params.useGPU then
      model:cuda()
   end
   local i = ffnn.FastTensor(input:narrow(1, 1, 1):size()):copy(input:narrow(1, 1, 1))
   return model:forward(i):nElement()
end

----------------------------------------------------------------------

function sizeForBatch(n, x)
   local size = x:size()
   size[1] = n
   return size
end

function fillBatch(data, first, batch, permutation)
   local actualBatchSize = math.min(params.batchSize, data.input:size(1) - first + 1)

   if batch.input then
      if actualBatchSize ~= batch.input:size(1) then
         batch.input:resize(sizeForBatch(actualBatchSize, batch.input))
      end
   else
      if torch.isTypeOf(data.input, ffnn.SlowTensor) then
         batch.input = ffnn.FastTensor(sizeForBatch(actualBatchSize, data.input));
      else
         batch.input = data.input.new():resize(sizeForBatch(actualBatchSize, data.input));
      end
   end

   if batch.target then
      if actualBatchSize ~= batch.target:size(1) then
         batch.target:resize(sizeForBatch(actualBatchSize, batch.target))
      end
   else
      if torch.isTypeOf(data.target, ffnn.SlowTensor) then
         batch.target = ffnn.FastTensor(sizeForBatch(actualBatchSize, data.target));
      else
         batch.target = data.target.new():resize(sizeForBatch(actualBatchSize, data.target));
      end
   end

   for k = 1, actualBatchSize do
      local i
      if permutation then
         i = permutation[first + k - 1]
      else
         i = first + k - 1
      end
      batch.input[k] = data.input[i]
      batch.target[k] = data.target[i]
   end
end

----------------------------------------------------------------------

--[[

The combineImage function takes as input a parameter c which is the
value to use for the background of the resulting image (padding and
such), and t which is either a 2d tensor, a 3d tensor, or a table.

 * If t is a 3d tensor, it is returned unchanged.

 * If t is a 2d tensor [r x c], it is reshaped to [1 x r x c] and
   returned.

 * If t is a table, combineImage first calls itself recursively on
   t[1], t[2], etc.

   It then creates a new tensor by concatenating the results
   horizontally if t.vertical is nil, vertically otherwise.

   It adds a padding of t.pad pixels if this field is set.

 * Example

   x = torch.Tensor(64, 64):fill(0.5)
   y = torch.Tensor(100, 30):fill(0.85)

   i = combineImages(1.0,
      {
         pad = 1,
         vertical = true,
         { pad = 1, x },
         {
            y,
            { pad = 4, torch.Tensor(32, 16):fill(0.25) },
            { pad = 1, torch.Tensor(45, 54):uniform(0.25, 0.9) },
         }
      }
   )

   image.save('example.png', i)

]]--

function combineImages(c, t)

   if torch.isTensor(t) then

      if t:dim() == 3 then
         return t
      elseif t:dim() == 2 then
         return torch.Tensor(1, t:size(1), t:size(2)):copy(t)
      else
         error('can only deal with [height x width] or [channel x height x width] tensors.')
      end

   else

      local subImages = {} -- The subimages
      local nc = 0 -- Nb of columns
      local nr = 0 -- Nb of rows

      for i, x in ipairs(t) do
         subImages[i] = combineImages(c, x)
         if t.vertical then
            nr = nr + subImages[i]:size(2)
            nc = math.max(nc, subImages[i]:size(3))
         else
            nr = math.max(nr, subImages[i]:size(2))
            nc = nc + subImages[i]:size(3)
         end
      end

      local pad = t.pad or 0
      local result = torch.Tensor(subImages[1]:size(1), nr + 2 * pad, nc + 2 * pad):fill(c)
      local co = 1 + pad -- Origin column
      local ro = 1 + pad -- Origin row

      for i in ipairs(t) do

         result
            :sub(1, subImages[1]:size(1),
                 ro, ro + subImages[i]:size(2) - 1,
                 co, co + subImages[i]:size(3) - 1)
            :copy(subImages[i])

         if t.vertical then
            ro = ro + subImages[i]:size(2)
         else
            co = co + subImages[i]:size(3)
         end

      end

      return result

   end

end

--[[

The imageFromTensors function gets as input a list of tensors of
arbitrary dimensions each, but whose two last dimensions stand for
height x width. It creates an image tensor (2d, one channel) with each
argument tensor unfolded per row.

]]--

function imageFromTensors(bt, signed)
   local gap = 1
   local tgap = -1
   local width = 0
   local height = gap

   for _, t in pairs(bt) do
      local d = t:dim()
      local h, w = t:size(d - 1), t:size(d)
      local n = t:nElement() / (w * h)
      width = math.max(width, gap + n * (gap + w))
      height = height + gap + tgap + gap + h
   end

   local e = torch.Tensor(3, height, width):fill(1.0)
   local y0 = 1 + gap

   for _, t in pairs(bt) do
      local d = t:dim()
      local h, w = t:size(d - 1), t:size(d)
      local n = t:nElement() / (w * h)
      local z = t:norm() / math.sqrt(t:nElement())

      local x0 = 1 + gap + math.floor( (width - n * (w + gap)) /2 )
      local u = torch.Tensor(t:size()):copy(t):resize(n, h, w)
      for m = 1, n do

         for c = 1, 3 do
            for y = 0, h+1 do
               e[c][y0 + y - 1][x0     - 1] = 0.0
               e[c][y0 + y - 1][x0 + w    ] = 0.0
            end
            for x = 0, w+1 do
               e[c][y0     - 1][x0 + x - 1] = 0.0
               e[c][y0 + h    ][x0 + x - 1] = 0.0
            end
         end

         for y = 1, h do
            for x = 1, w do
               local v = u[m][y][x] / z
               local r, g, b
               if signed then
                  if v < -1 then
                     r, g, b = 0.0, 0.0, 1.0
                  elseif v > 1 then
                     r, g, b = 1.0, 0.0, 0.0
                  elseif v >= 0 then
                     r, g, b = 1.0, 1.0 - v, 1.0 - v
                  else
                     r, g, b = 1.0 + v, 1.0 + v, 1.0
                  end
               else
                  if v <= 0 then
                     r, g, b = 1.0, 1.0, 1.0
                  elseif v > 1 then
                     r, g, b = 0.0, 0.0, 0.0
                  else
                     r, g, b = 1.0 - v, 1.0 - v, 1.0 - v
                  end
               end
               e[1][y0 + y - 1][x0 + x - 1] = r
               e[2][y0 + y - 1][x0 + x - 1] = g
               e[3][y0 + y - 1][x0 + x - 1] = b
            end
         end
         x0 = x0 + w + gap
      end
      y0 = y0 + h + gap + tgap + gap
   end

   return e
end