import os
from pycocotools.coco import COCO
from PIL import Image
import numpy as np

cocoDir = '/disk3/lilong/coco_2014'
dataType = ['train2014', 'val2014']

fixationDir = '/disk3/lilong/coco_2014/fixation_maps_org'

outputDir = '/disk3/lilong/rank_saliancy'

annFiles = ['{}/annotations/instances_{}.json'.format(cocoDir, dType) for dType in dataType]
cocos = [COCO(annFile) for annFile in annFiles]
coco_img_names_list = [[img['file_name'] for img in list(coco.imgs.values())] for coco in cocos]

fix_lists = os.listdir(fixationDir)
fix_imgs = dict()
for fix in fix_lists:
    coco = cocos[0] if fix.split('_')[1] == 'train2014' else cocos[1]
    coco_img_names = coco_img_names_list[0] if fix.split('_')[1] == 'train2014' else coco_img_names_list[1]
    fix_img_name = '{}.jpg'.format(fix[:-4])
    img_idx = list(coco.imgs.keys())[coco_img_names.index(fix_img_name)]
    fix_imgs[img_idx] = {}
    fix_imgs[img_idx]['type'] = fix.split('_')[1]
    fix_imgs[img_idx]['img_name'] = fix_img_name

idxs = list(sorted(fix_imgs.keys()))
for idx in idxs:
    coco = cocos[0] if fix_imgs[idx]['type'] == 'train2014' else cocos[1]

    fix_img_name = '{}.png'.format(fix_imgs[idx]['img_name'][:-4])
    fix_img_path = os.path.join(fixationDir, fix_img_name)
    fix_image = Image.open(fix_img_path)
    fix_image = np.asarray(fix_image)

    coco_image = coco.imgs[idx]
    coco_imageSize = (coco_image['height'], coco_image['width'])
    labelMap = np.zeros(coco_imageSize)
    check_mask = np.zeros(coco_imageSize)
    imgAnnots = [a for a in coco.anns.values() if a['image_id'] == idx]
    imgAnnots = sorted(imgAnnots, key=lambda x: x['area'])
    instances = {}
    for a in range(0, len(imgAnnots)):
        instances[a] = {}
        mask = coco.annToMask(imgAnnots[a])
        labelMap_mask = check_mask != 0
        mask[labelMap_mask] = 0
        labelMask = mask == 1
        check_mask[labelMask] = a+1
        fix_density = (fix_image[labelMask] == 255).sum()
        instances[a]['fix_density'] = fix_density
        instances[a]['labelMask'] = labelMask
    instances = sorted(instances.items(), key=lambda x: x[1]['fix_density'])
    for i in range(0, len(instances)):
        rank = int(255./len(instances)*(i+1))
        labelMap[instances[i][1]['labelMask']] = rank
    png = Image.fromarray(labelMap).convert('P')
    out_dir = os.path.join(outputDir, fix_img_name)
    png.save(out_dir, format='PNG')