import numpy as np
import pickle
from pdb import set_trace as bp

cls_true = np.load('store_cls.npy')
scores_bb = np.load('store_scores.npy')
pickleFile = open('store_ious_cls.pkl',"rb")
ious_cls = pickle.load(pickleFile)
pickleFile.close()
ious_cls = [ious_cls[i] for i in range(1,21)]

all_mean = scores_bb.mean()

# all_sample_ious = []
# store_cls_means = []
# for i, iou_c in enumerate(ious_cls):
#     iou_c = np.array(iou_c)
#     scores_c = scores_bb[cls_true==i+1]
#     assert iou_c.shape==scores_c.shape
#     iou_c_filter = iou_c[scores_c>all_mean]
#     print(f'cls: {i+1}, n_earlier: {iou_c.shape[0]}, n_now: {iou_c_filter.shape[0]}, IoU: {iou_c_filter.mean()}')
#     store_cls_means.append(iou_c_filter.mean())
#     all_sample_ious+=list(iou_c_filter)
# print(f'all cls mean: {np.mean(store_cls_means)}, all sample mean: {np.mean(all_sample_ious)}')

store_ious = np.load('store_ious.npy')
all_sample_ious = []
store_cls_means = []
store_cls_stds = []
for i in range(20):
    iou_c = store_ious[cls_true==i+1]
    scores_c = scores_bb[cls_true==i+1]
    assert iou_c.shape==scores_c.shape
    # iou_c_filter = iou_c[scores_c>1.05*all_mean]
    # iou_c_filter = iou_c[scores_c>scores_c.mean()+0.25*scores_c.std()]
    iou_c_filter = iou_c[np.argpartition(scores_c, -int(0.52*scores_c.shape[0]))[-int(0.52*scores_c.shape[0]):]]
    print(f'cls: {i+1}, n_earlier: {iou_c.shape[0]}, n_now: {iou_c_filter.shape[0]}, IoU: {iou_c_filter.mean()}')
    store_cls_means.append(iou_c_filter.mean())
    store_cls_stds.append(iou_c_filter.std())
    all_sample_ious+=list(iou_c_filter)
bp()
print(f'all cls mean: {np.mean(store_cls_means)} and std {np.mean(store_cls_stds)}, all sample mean: {np.mean(all_sample_ious)}')