--[[ An auRoc class

-- Same form as the torch optim.ConfusionMatrix class

-- output is assumed to be a ranking (e.g. probability value)

-- label is assumed to be 1 or -1

Example:

    auRoc = auRoc.new()   -- new matrix

    conf:zero()                                              -- reset matrix

    for i = 1,N do

        conf:add( output, label )         -- accumulate errors

    end

    print(auRoc:calculateAuc())

]]

local auRoc = torch.class("auRoc")

function auRoc:__init()

   self.target = {}

   self.pred = {}

   self.roc = 0

   self.auc = 0

end

function auRoc:add(prediction, target)

  if target == 2 or target == 0 or target == -1 then

    target = -1

  elseif target == 1 then

    target = 1

  else

    print('Incorrect target for auRoc:add(). Exiting')

    os.exit()

  end

  table.insert(self.pred,prediction)

  table.insert(self.target,target)

end

function auRoc:zero()

   self.target = {}

   self.pred = {}

   self.roc = 0

   self.auc = 0

end

local function tableToTensor(table)

  local tensor = torch.Tensor(#table)

  for i = 1,#table do

    tensor[i] = table[i]

  end

  return tensor

end

local function get_rates(responses, labels)

  torch.setdefaulttensortype('torch.FloatTensor')

  responses = torch.Tensor(responses:size()):copy(responses)

  labels = torch.Tensor(labels:size()):copy(labels)

   -- assertions about the data format expected

   assert(responses:size():size() == 1, "responses should be a 1D vector")

   assert(labels:size():size() == 1 , "labels should be a 1D vector")

   -- assuming labels {-1, 1}

   local npositives = torch.sum(torch.eq(labels,  1))

   local nnegatives = torch.sum(torch.eq(labels, -1))

   local nsamples = npositives + nnegatives

   assert(nsamples == responses:size()[1], "labels should contain only -1 or 1 values")

   -- sort by response value

   local responses_sorted, indexes_sorted = torch.sort(responses,1,true)

   local labels_sorted = labels:index(1, indexes_sorted)

   local found_positives = 0

   local found_negatives = 0

   local tpr = {0} -- true pos rate

   local fpr = {0} -- false pos rate

   for i = 1,nsamples-1 do

      if labels_sorted[i] == -1 then

         found_negatives = found_negatives + 1

      else

         found_positives = found_positives + 1

      end

      table.insert(tpr, found_positives/npositives)

      table.insert(fpr, found_negatives/nnegatives)

   end

   table.insert(tpr, 1.0)

   table.insert(fpr, 1.0)

   return tpr, fpr

end

local function find_auc(tpr,fpr)

   local area = 0.0

   for i = 2,#tpr do

      local xdiff = fpr[i] - fpr[i-1]

      local ydiff = tpr[i] - tpr[i-1]

      area = area + (xdiff * tpr[i])

   end

   return area

end

function auRoc:calculateAuc()

  local aucPredTens = tableToTensor(self.pred)

  local aucTargetTens = tableToTensor(self.target)

  local tpr = nil

  local fpr = nil

  tpr,fpr = get_rates(aucPredTens,aucTargetTens)

  self.auc = find_auc(tpr,fpr)

  return self.auc

end