--[[----------------------------------------------------------------------------

Interface for accessing the Common Objects in COntext (COCO) dataset.

For an overview of the API please see http://mscoco.org/dataset/#download.
CocoApi.lua (this file) is modeled after the Matlab CocoApi.m:
  https://github.com/pdollar/coco/blob/master/MatlabAPI/CocoApi.m

The following API functions are defined in the Lua API:
  CocoApi    - Load COCO annotation file and prepare data structures.
  getAnnIds  - Get ann ids that satisfy given filter conditions.
  getCatIds  - Get cat ids that satisfy given filter conditions.
  getImgIds  - Get img ids that satisfy given filter conditions.
  loadAnns   - Load anns with the specified ids.
  loadCats   - Load cats with the specified ids.
  loadImgs   - Load imgs with the specified ids.
  showAnns   - Display the specified annotations.
Throughout the API "ann"=annotation, "cat"=category, and "img"=image.
For detailed usage information please see cocoDemo.lua.

LIMITATIONS: the following API functions are NOT defined in the Lua API:
  loadRes    - Load algorithm results and create API for accessing them.
  download   - Download COCO images from mscoco.org server.
In addition, currently the getCatIds() and getImgIds() do not accept filters.
getAnnIds() can be called using getAnnIds({imgId=id}) and getAnnIds({catId=id}).

Note: loading COCO JSON annotations to Lua tables is quite slow. Hence, a call
to CocApi(annFile) converts the annotations to a custom 'flattened' format that
is more efficient. The first time a COCO JSON is loaded, the conversion is
invoked (this may take up to a minute). The converted data is then stored in a
t7 file (the code must have write permission to the dir of the JSON file).
Future calls of cocoApi=CocApi(annFile) take a fraction of a second. To view the
created data just inspect cocoApi.data of a created instance of the CocoApi.

Common Objects in COntext (COCO) Toolbox.      version 3.0
Data, paper, and tutorials available at:  http://mscoco.org/
Code written by Pedro O. Pinheiro and Piotr Dollar, 2016.
Licensed under the Simplified BSD License [see coco/license.txt]

------------------------------------------------------------------------------]]

local json = require 'cjson'
local coco = require 'coco.env'

local TensorTable = torch.class('TensorTable',coco)
local CocoSeg = torch.class('CocoSeg',coco)
local CocoApi = torch.class('CocoApi',coco)

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

--[[ TensorTable is a lightweight data structure for storing variable size 1D
tensors. Tables of tensors are slow to save/load to disk. Instead, TensorTable
stores all the data in a single long tensor (along with indices into the tensor)
making serialization fast. A TensorTable may only contain 1D same-type torch
tensors or strings. It supports only creation from a table and indexing. ]]

function TensorTable:__init( T )
  local n = #T; assert(n>0)
  local isStr = torch.type(T[1])=='string'
  assert(isStr or torch.isTensor(T[1]))
  local c=function(s) return torch.CharTensor(torch.CharStorage():string(s)) end
  if isStr then local S=T; T={}; for i=1,n do T[i]=c(S[i]) end end
  local ms, idx = torch.LongTensor(n), torch.LongTensor(n+1)
  for i=1,n do ms[i]=T[i]:numel() end
  idx[1]=1; idx:narrow(1,2,n):copy(ms); idx=idx:cumsum()
  local type = string.sub(torch.type(T[1]),7,-1)
  local data = torch[type](idx[n+1]-1)
  if isStr then type='string' end
  for i=1,n do if ms[i]>0 then data:sub(idx[i],idx[i+1]-1):copy(T[i]) end end
  if ms:eq(ms[1]):all() and ms[1]>0 then data=data:view(n,ms[1]); idx=nil end
  self.data, self.idx, self.type = data, idx, type
end

function TensorTable:__index__( i )
  if torch.type(i)~='number' then return false end
  local d, idx, type = self.data, self.idx, self.type
  if idx and idx[i]==idx[i+1] then
    if type=='string' then d='' else d=torch[type]() end
  else
    if idx then d=d:sub(idx[i],idx[i+1]-1) else d=d[i] end
    if type=='string' then d=d:clone():storage():string() end
  end
  return d, true
end

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

--[[ CocoSeg is an efficient data structure for storing COCO segmentations. ]]

function CocoSeg:__init( segs )
  local polys, pIdx, sizes, rles, p, isStr = {}, {}, {}, {}, 0, 0
  for i,seg in pairs(segs) do if seg.size then isStr=seg.counts break end end
  isStr = torch.type(isStr)=='string'
  for i,seg in pairs(segs) do
    pIdx[i], sizes[i] = {}, {}
    if seg.size then
      sizes[i],rles[i] = seg.size,seg.counts
    else
      if isStr then rles[i]='' else rles[i]={} end
      for j=1,#seg do p=p+1; pIdx[i][j],polys[p] = p,seg[j] end
    end
    pIdx[i],sizes[i] = torch.LongTensor(pIdx[i]),torch.IntTensor(sizes[i])
    if not isStr then rles[i]=torch.IntTensor(rles[i]) end
  end
  for i=1,p do polys[i]=torch.DoubleTensor(polys[i]) end
  self.polys, self.pIdx = coco.TensorTable(polys), coco.TensorTable(pIdx)
  self.sizes, self.rles = coco.TensorTable(sizes), coco.TensorTable(rles)
end

function CocoSeg:__index__( i )
  if torch.type(i)~='number' then return false end
  if self.sizes[i]:numel()>0 then
    return {size=self.sizes[i],counts=self.rles[i]}, true
  else
    local ids, polys = self.pIdx[i], {}
    for i=1,ids:numel() do polys[i]=self.polys[ids[i]] end
    return polys, true
  end
end

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

--[[ CocoApi is the API to the COCO dataset, see main comment for details. ]]

function CocoApi:__init( annFile )
  assert( string.sub(annFile,-4,-1)=='json' and paths.filep(annFile) )
  local torchFile = string.sub(annFile,1,-6) .. '.t7'
  if not paths.filep(torchFile) then self:__convert(annFile,torchFile) end
  local data = torch.load(torchFile)
  self.data, self.inds = data, {}
  for k,v in pairs({images='img',categories='cat',annotations='ann'}) do
    local M = {}; self.inds[v..'IdsMap']=M
    if data[k] then for i=1,data[k].id:size(1) do M[data[k].id[i]]=i end end
  end
end

function CocoApi:__convert( annFile, torchFile )
  print('convert: '..annFile..' --> .t7 [please be patient]')
  local tic = torch.tic()
  -- load data and decode json
  local data = torch.CharStorage(annFile):string()
  data = json.decode(data); collectgarbage()
  -- transpose and flatten each field in the coco data struct
  local convert = {images=true, categories=true, annotations=true}
  for field, d in pairs(data) do if convert[field] then
    print('converting: '..field)
    local n, out = #d, {}
    if n==0 then d,n={d},1 end
    for k,v in pairs(d[1]) do
      local t, isReg = torch.type(v), true
      for i=1,n do isReg=isReg and torch.type(d[i][k])==t end
      if t=='number' and isReg then
        out[k] = torch.DoubleTensor(n)
        for i=1,n do out[k][i]=d[i][k] end
      elseif t=='string' and isReg then
        out[k]={}; for i=1,n do out[k][i]=d[i][k] end
        out[k] = coco.TensorTable(out[k])
      elseif t=='table' and isReg and torch.type(v[1])=='number' then
        out[k]={}; for i=1,n do out[k][i]=torch.DoubleTensor(d[i][k]) end
        out[k] = coco.TensorTable(out[k])
        if not out[k].idx then out[k]=out[k].data end
      else
        out[k]={}; for i=1,n do out[k][i]=d[i][k] end
        if k=='segmentation' then out[k] = coco.CocoSeg(out[k]) end
      end
      collectgarbage()
    end
    if out.id then out.idx=torch.range(1,out.id:size(1)) end
    data[field] = out
    collectgarbage()
  end end
  -- create mapping from cat/img index to anns indices for that cat/img
  print('convert: building indices')
  local makeMap = function( type, type_id )
    if not data[type] or not data.annotations then return nil end
    local invmap, n = {}, data[type].id:size(1)
    for i=1,n do invmap[data[type].id[i]]=i end
    local map = {}; for i=1,n do map[i]={} end
    data.annotations[type_id..'x'] = data.annotations[type_id]:clone()
    for i=1,data.annotations.id:size(1) do
      local id = invmap[data.annotations[type_id][i]]
      data.annotations[type_id..'x'][i] = id
      table.insert(map[id],data.annotations.id[i])
    end
    for i=1,n do map[i]=torch.LongTensor(map[i]) end
    return coco.TensorTable(map)
  end
  data.annIdsPerImg = makeMap('images','image_id')
  data.annIdsPerCat = makeMap('categories','category_id')
  -- save to disk
  torch.save( torchFile, data )
  print(('convert: complete [%.2f s]'):format(torch.toc(tic)))
end

function CocoApi:getAnnIds( filters )
  if not filters then filters = {} end
  if filters.imgId then
    return self.data.annIdsPerImg[self.inds.imgIdsMap[filters.imgId]] or {}
  elseif filters.catId then
    return self.data.annIdsPerCat[self.inds.catIdsMap[filters.catId]] or {}
  else
    return self.data.annotations.id
  end
end

function CocoApi:getCatIds()
  return self.data.categories.id
end

function CocoApi:getImgIds()
  return self.data.images.id
end

function CocoApi:loadAnns( ids )
  return self:__load(self.data.annotations,self.inds.annIdsMap,ids)
end

function CocoApi:loadCats( ids )
  return self:__load(self.data.categories,self.inds.catIdsMap,ids)
end

function CocoApi:loadImgs( ids )
  return self:__load(self.data.images,self.inds.imgIdsMap,ids)
end

function CocoApi:showAnns( img, anns )
  local n, h, w = #anns, img:size(2), img:size(3)
  local MaskApi, clrs = coco.MaskApi, torch.rand(n,3)*.6+.4
  local O = img:clone():contiguous():float()
  if n==0 then anns,n={anns},1 end
  if anns[1].keypoints then for i=1,n do if anns[i].iscrowd==0 then
    local sk, kp, j, k = self:loadCats(anns[i].category_id)[1].skeleton
    kp=anns[i].keypoints; k=kp:size(1); j=torch.range(1,k,3):long(); k=k/3;
    local x,y,v = kp:index(1,j), kp:index(1,j+1), kp:index(1,j+2)
    for _,s in pairs(sk) do if v[s[1]]>0 and v[s[2]]>0 then
      MaskApi.drawLine(O,x[s[1]],y[s[1]],x[s[2]],y[s[2]],.75,clrs[i])
    end end
    for j=1,k do if v[j]==1 then MaskApi.drawCirc(O,x[j],y[j],4,{0,0,0}) end end
    for j=1,k do if v[j]>0 then MaskApi.drawCirc(O,x[j],y[j],3,clrs[i]) end end
  end end end
  if anns[1].segmentation or anns[1].bbox then
    local Rs, alpha = {}, anns[1].keypoints and .25 or .4
    for i=1,n do
      Rs[i]=anns[i].segmentation
      if Rs[i] and #Rs[i]>0 then Rs[i]=MaskApi.frPoly(Rs[i],h,w) end
      if not Rs[i] then Rs[i]=MaskApi.frBbox(anns[i].bbox,h,w)[1] end
    end
    MaskApi.drawMasks(O,MaskApi.decode(Rs),nil,alpha,clrs)
  end
  return O
end

function CocoApi:__load( data, map, ids )
  if not torch.isTensor(ids) then ids=torch.LongTensor({ids}) end
  local out, idx = {}, nil
  for i=1,ids:numel() do
    out[i], idx = {}, map[ids[i]]
    for k,v in pairs(data) do out[i][k]=v[idx] end
  end
  return out
end