import sys 
import os 
import numpy as np
import torch 


def get_topn_same_recall(score, labels, N, threshold = 0.0):
    gt_num = list(labels).count(0)
    hit = 0.0

    for i in range(min(N, len(labels))):
        if labels[i] == 0 and score[i] >= threshold:
            hit += 1
    if gt_num == 0 or N == 0:
        return 0.0

    return hit / min(N, gt_num)



def get_topn_hit_rate(labels, N):
    """HR@N: Top-N中是否存在至少一个正样本(label==0)，是则为1，否则为0"""
    for i in range(min(N, len(labels))):
        if labels[i] == 0:
            return 1.0
    return 0.0


def TopNHitRate(query_dic, TopN):
    """计算所有query的平均HR@N"""
    hr_sum = 0.0
    query_num = 0.0

    for _, data in query_dic.items():
        query_score = np.array(data['score'])
        query_label = np.array(data['label'])

        index_score = np.argsort(-query_score)
        query_label_sort = query_label[index_score]

        if 0 in query_label:
            hr_sum += get_topn_hit_rate(query_label_sort, TopN)
            query_num += 1

    if query_num == 0:
        return 0.0
    return hr_sum / query_num


def TopNAcc(query_dic, TopN):
    same_ap = 0.0
    same_query_num = 0.0
    
    for _, data in query_dic.items():
        query_score = np.array(data['score'])
        query_label = np.array(data['label'])

        index_score = np.argsort(-query_score)
        query_score_sort = query_score[index_score]
        query_label_sort = query_label[index_score]

        if 0 in query_label:
            same_ap += get_topn_same_recall(query_score_sort, query_label_sort, TopN)
            same_query_num += 1
    
    if same_query_num == 0:
        return 0.0
    return same_ap / same_query_num


def tongkuan_score_print(query_dic_all, sim_weight):
    results_all = []
    results_all.append(['same_Top1','same_Top5','same_Top10', 'same_Top20','same_Top30','sim', 'HR@1', 'HR@5', 'HR@10', 'HR@20'])

    query_score_dic = {}
    for query, value in query_dic_all.items():
        query_score_dic[query] = {}
        query_score_dic[query]['score'] = value['sim']
        query_score_dic[query]['label'] = value['label']       

    # 基于上面的数据，去算相似率
    same_rate_Top1 = TopNAcc(query_score_dic, 1)
    same_rate_Top5 = TopNAcc(query_score_dic, 5)
    same_rate_Top10 = TopNAcc(query_score_dic, 10)
    same_rate_Top20 = TopNAcc(query_score_dic, 20)
    same_rate_Top30 = TopNAcc(query_score_dic, 30)
    sim = sim_weight[0] * same_rate_Top1 + sim_weight[1] * same_rate_Top5 + sim_weight[2] * same_rate_Top10 + sim_weight[3] * same_rate_Top20 + sim_weight[4] * same_rate_Top30

    hr_1 = TopNHitRate(query_score_dic, 1)
    hr_5 = TopNHitRate(query_score_dic, 5)
    hr_10 = TopNHitRate(query_score_dic, 10)
    hr_20 = TopNHitRate(query_score_dic, 20)

    print ('##### wo margin #####')
    print ('same_rate_Top1: ', same_rate_Top1)
    print ('same_rate_Top5: ', same_rate_Top5)
    print ('same_rate_Top10: ', same_rate_Top10)
    print ('same_rate_Top20: ', same_rate_Top20)
    print ('same_rate_Top30: ', same_rate_Top30)
    print ('sim: ', sim)
    print ('HR@1: ', hr_1)
    print ('HR@5: ', hr_5)
    print ('HR@10: ', hr_10)
    print ('HR@20: ', hr_20)

    results_all.append([same_rate_Top1, same_rate_Top5, same_rate_Top10, same_rate_Top20, same_rate_Top30, sim, hr_1, hr_5, hr_10, hr_20])

    results_all_dict = {}
    for i in range(len(results_all[0])):
        results_all_dict[results_all[0][i]] = results_all[1][i]
    return results_all_dict