from scipy.stats import spearmanr
import sys
import numpy as np


# Similarity benchmark evaluation
def eval_queryset(vsm, queryset, pos_filter=None, verbose=False):
    gold_scores = []
    predicted_scores = []
    for w1, w2, pos in queryset.queries:
        if pos_filter is None or pos_filter == pos:
            gold_scores += [queryset.queries[(w1, w2, pos)]]
            w1, w2 = w1.lower(), w2.lower()
            predicted_score = vsm.similarity(w1, w2, pos, pos)
            predicted_scores += [predicted_score]

    na, total, _ = queryset.coverage(vsm)
    if verbose:
        sys.stderr.write("Unknown words in queries: {} / {}\n".format(na, total))
    return spearmanr(gold_scores, predicted_scores).correlation


# Word class suggestion evaluation for a single predicate
def eval_class_suggestion(predicted_classes, true_classes, max_k=10):

    # rank predictions: classes with same score get grouped into the same rank
    def rank(classes_and_scores):
        rx = {}
        for (a, b, _) in classes_and_scores:
            rx[b] = rx.get(b, []) + [a]
        vals = sorted(rx.keys(), reverse=True)
        rx_ranked = {i + 1: rx[vals[i]] for i in range(len(vals))}
        return rx_ranked

    # collect predictions up to rank k
    def inclusive_get(classes_and_scores_r, max_rank):
        res = []
        for r in range(max_rank + 1):
            res += classes_and_scores_r.get(r, [])
        return res

    # calculate recall for different values of k
    def R_score(predicted_classes, true_classes, max_k):
        R = {}
        predictions_r = rank(predicted_classes)
        last_r = 0.0
        for k in sorted(predictions_r.keys()):
            intersection = set(inclusive_get(predictions_r, k)).intersection(set(true_classes))
            R[k] = len(intersection)*1.0/len(true_classes)
            last_r = R[k]
            if k >= max_k:
                break

        for k in range(1, max_k+1):
            if k not in R:  # if no items on rank k (e.g. for prototype)
                R[k] = last_r
        return R

    # calculate precision for different values of k
    def P_score(predicted_classes, true_classes, max_k):
        P = {}
        predictions_p = rank(predicted_classes)
        last_p = 0.0
        for k in sorted(predictions_p.keys()):
            intersection = set(inclusive_get(predictions_p, k)).intersection(set(true_classes))
            P[k] = len(intersection)*1.0/len(inclusive_get(predictions_p, k))
            last_p = P[k]
            if k >= max_k:
                break
        for k in range(1, max_k + 1):
            if k not in P:
                P[k] = last_p
        return P

    # calculate F from P and R for different values of k
    def F_score(r, p):
        F = {}
        for k in r:
            F[k] = (2.0 * p[k] * r[k] / (p[k] + r[k])) if (p[k] + r[k] != 0.0) else 0.0
        return F

    def vec(voc, max_k):
        arr = np.zeros(max_k)
        for i in voc:
            arr[i - 1] = voc[i]
        return arr

    P = P_score(predicted_classes, true_classes, max_k)
    R = R_score(predicted_classes, true_classes, max_k)

    return {"P": vec(P, max_k),
            "R": vec(R, max_k),
            "F": vec(F_score(R, P), max_k)}