import numpy as np
import json
import pickle
import argparse
import mmcv

num_classes = 8

def parse_args():
    parser = argparse.ArgumentParser(
            description='filtering results')
    parser.add_argument('rname')
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()
    rname = args.rname
    
    gt = mmcv.load('/data/dibe/data/test.json')
    class_num = [0] * num_classes
    for i in range(num_classes):
        class_num[gt['categories'][i]['id']-1] = i

    gtval = mmcv.load('/data/dibe/data/train.json')
    gt_num = [0] * num_classes
    for i in gtval['annotations']:
        gt_num[class_num[i['category_id']-1]] += 1
    gt_num = np.array(gt_num)
    gt_num = gt_num / len(gtval['images'])
    
    print(len(gtval['annotations']), len(gtval['images']))
    
    aaa = mmcv.load(rname.split('.')[0] + '.bbox.json')
    
    bc = [[] for i in range(num_classes)]
    for bbox in aaa:
        bc[class_num[bbox['category_id']-1]].append(bbox)
    
    
    thresh_list = np.linspace(0.1, 1.0, 91)
    thresh = np.zeros((num_classes))
    
    # seperate
    for nc in range(num_classes):
        r = np.zeros(len(thresh_list))
        for t in range(len(thresh_list)):
            ind = [i for i in bc[nc] if i['score'] > thresh_list[t]] 
            r[t] = len(ind)/len(gt['images'])

        ii = np.argmin(np.abs(r-gt_num[nc]))
        thresh[nc] = thresh_list[ii]
        print(nc, thresh[nc])
    
    
    cci = [i['id'] for i in gt['categories']]
    c = []
    
    
    for i in range(len(aaa)):
        if aaa[i]['score'] > thresh[cci.index(aaa[i]['category_id'])]:
            c.append(aaa[i])
    
    print(len(c))
    
    json.dump(c, open(rname.split('.')[0] + '_t.bbox.json','w'))
    
    a = pickle.load(open(rname, 'rb'))
    c = a.copy()
    
    for i in range(0, len(a)):
        for nc in range(num_classes):
            ind = np.where(a[i][nc][:,-1]>thresh[nc])[0]
            c[i][nc] = a[i][nc][ind]
    
    print(sum([np.vstack(i).shape[0] for i in c]))
    import pickle
    pickle.dump(c, open(rname.split('.')[0] + '_t.pkl','wb'))