"""
This module provides some utils for calculating metrics
"""
import numpy as np
from sklearn.metrics import average_precision_score


def softmax(raw_score, T=1):
    exp_s = np.exp((raw_score - raw_score.max(axis=-1)[..., None])*T)
    sum_s = exp_s.sum(axis=-1)
    return exp_s / sum_s[..., None]


def top_k_acc(lb_set, scores, k=3):
    idx = np.argsort(scores)[-k:]
    return len(lb_set.intersection(idx)), len(lb_set)


def top_k_hit(lb_set, scores, k=3):
    idx = np.argsort(scores)[-k:]
    return len(lb_set.intersection(idx)) > 0, 1


def top_3_accuracy(score_dict, video_list):
    return top_k_accuracy(score_dict, video_list, 3)


def top_k_accuracy(score_dict, video_list, k):
    video_labels = [set([i.num_label for i in v.instances]) for v in video_list]

    video_top_k_acc = np.array(
        [top_k_hit(lb, score_dict[v.id], k=k) for v, lb in zip(video_list, video_labels)
         if v.id in score_dict])

    tmp = video_top_k_acc.sum(axis=0).astype(float)
    top_k_acc = tmp[0] / tmp[1]

    return top_k_acc


def video_mean_ap(score_dict, video_list):
    avail_video_labels = [set([i.num_label for i in v.instances]) for v in video_list if
                          v.id in score_dict]
    pred_array = np.array([score_dict[v.id] for v in video_list if v.id in score_dict])
    gt_array = np.zeros(pred_array.shape)

    for i in xrange(pred_array.shape[0]):
        gt_array[i, list(avail_video_labels[i])] = 1
    mean_ap = average_precision_score(gt_array, pred_array, average='macro')
    return mean_ap
