from common import Method
from common.SaveLoadPOPO import SaveLoadPOPO
from util.save_load import SaveLoad
import numpy as np
from model import svm

import scipy.sparse as sp
class GetRankings(Method.Method):
    word_dir = None
    words = None
    output_directions = None
    output_folder = None
    directions = None
    dirs = None
    bowmin = None
    bowmax = None
    new_word_dict = None
    space = None
    LR = None

    def __init__(self, rankings,  words, save_class, file_name, output_folder, bowmin, bowmax):
        self.rankings = rankings
        self.output_folder = output_folder
        super().__init__(file_name, save_class)

    def makePopos(self):

        self.ndcg = SaveLoadPOPO(np.empty(len(self.rankings), dtype=object), self.output_folder + "rank/" + self.file_name + "_" + str(self.bowmin) + "_" + str(self.bowmax) + "_rank.npy", ".npy")


    def makePopoArray(self):
        self.popo_array = [self.rankings, self.rank_dir]

    def get_dp(self, direction):
        ranking = np.empty(len(self.space))
        for i in range(len(self.space)):
            ranking[i] = np.dot(direction, self.space[i])
        return ranking

    def process(self):
        for i in range(len(self.words)):
            if self.words[i] not in self.rank_dir.value:
                self.rank_dir.value[self.words[i]] = self.get_dp(self.dirs[i])
            self.rankings.value[i] = self.rank_dir.value(self.words[i])
        super().process()

    def getRankings(self):
        if self.processed is False:
            self.rankings.value = self.save_class.load(self.rankings)
        return self.rankings.value


def ranking_precision_score(y_true, y_score, k=10):
    """Precision at rank k
    Parameters
    ----------
    y_true : array-like, shape = [n_samples]
        Ground truth (true relevance labels).
    y_score : array-like, shape = [n_samples]
        Predicted scores.
    k : int
        Rank.
    Returns
    -------
    precision @k : float
    """
    unique_y = np.unique(y_true)

    if len(unique_y) > 2:
        raise ValueError("Only supported for two relevance levels.")

    pos_label = unique_y[1]
    n_pos = np.sum(y_true == pos_label)

    order = np.argsort(y_score)[::-1]
    y_true = np.take(y_true, order[:k])
    n_relevant = np.sum(y_true == pos_label)

    # Divide by min(n_pos, k) such that the best achievable score is always 1.0.
    return float(n_relevant) / min(n_pos, k)


def average_precision_score(y_true, y_score, k=10):
    """Average precision at rank k
    Parameters
    ----------
    y_true : array-like, shape = [n_samples]
        Ground truth (true relevance labels).
    y_score : array-like, shape = [n_samples]
        Predicted scores.
    k : int
        Rank.
    Returns
    -------
    average precision @k : float
    """
    unique_y = np.unique(y_true)

    if len(unique_y) > 2:
        raise ValueError("Only supported for two relevance levels.")

    pos_label = unique_y[1]
    n_pos = np.sum(y_true == pos_label)

    order = np.argsort(y_score)[::-1][:min(n_pos, k)]
    y_true = np.asarray(y_true)[order]

    score = 0
    for i in range(len(y_true)):
        if y_true[i] == pos_label:
            # Compute precision up to document i
            # i.e, percentage of relevant documents up to document i.
            prec = 0
            for j in range(0, i + 1):
                if y_true[j] == pos_label:
                    prec += 1.0
            prec /= (i + 1.0)
            score += prec

    if n_pos == 0:
        return 0

    return score / n_pos


def dcg_score(y_true, y_score, k=10, gains="exponential"):
    """Discounted cumulative gain (DCG) at rank k
    Parameters
    ----------
    y_true : array-like, shape = [n_samples]
        Ground truth (true relevance labels).
    y_score : array-like, shape = [n_samples]
        Predicted scores.
    k : int
        Rank.
    gains : str
        Whether gains should be "exponential" (default) or "linear".
    Returns
    -------
    DCG @k : float
    """
    order = np.argsort(y_score)[::-1]
    y_true = np.take(y_true, order[:k])

    if gains == "exponential":
        gains = 2 ** y_true - 1
    elif gains == "linear":
        gains = y_true
    else:
        raise ValueError("Invalid gains option.")

    # highest rank is 1 so +2 instead of +1
    discounts = np.log2(np.arange(len(y_true)) + 2)
    return np.sum(gains / discounts)


def ndcg_score(y_true, y_score, k=10, gains="exponential"):
    """Normalized discounted cumulative gain (NDCG) at rank k
    Parameters
    ----------
    y_true : array-like, shape = [n_samples]
        Ground truth (true relevance labels).
    y_score : array-like, shape = [n_samples]
        Predicted scores.
    k : int
        Rank.
    gains : str
        Whether gains should be "exponential" (default) or "linear".
    Returns
    -------
    NDCG @k : float
    """
    best = dcg_score(y_true, y_true, k, gains)
    actual = dcg_score(y_true, y_score, k, gains)
    return actual / best


# Alternative API.

def dcg_from_ranking(y_true, ranking):
    """Discounted cumulative gain (DCG) at rank k
    Parameters
    ----------
    y_true : array-like, shape = [n_samples]
        Ground truth (true relevance labels).
    ranking : array-like, shape = [k]
        Document indices, i.e.,
            ranking[0] is the index of top-ranked document,
            ranking[1] is the index of second-ranked document,
            ...
    k : int
        Rank.
    Returns
    -------
    DCG @k : float
    """
    y_true = np.asarray(y_true)
    ranking = np.asarray(ranking)
    rel = y_true[ranking]
    gains = 2 ** rel - 1
    discounts = np.log2(np.arange(len(ranking)) + 2)
    overall = gains / discounts
    sum = np.sum(overall)
    return sum


def ndcg_from_ranking(y_true, ranking):
    """Normalized discounted cumulative gain (NDCG) at rank k
    Parameters
    ----------
    y_true : array-like, shape = [n_samples]
        Ground truth (true relevance labels).
    ranking : array-like, shape = [k]
        Document indices, i.e.,
            ranking[0] is the index of top-ranked document,
            ranking[1] is the index of second-ranked document,
            ...
    k : int
        Rank.
    Returns
    -------
    NDCG @k : float
    """
    k = len(ranking)
    best_ranking = np.argsort(y_true)[::-1]
    best = dcg_from_ranking(y_true, best_ranking[:k])
    dcg = dcg_from_ranking(y_true, ranking)
    return dcg / best


def getNDCG(rankings_fn, fn, data_type, bow_fn, ppmi_fn, lowest_count, rewrite_files=False, highest_count = 0, classification = ""):

    # Check if the NDCG scores have already been calculated, if they have then skip.


    # Get the file names for the PPMI values for every word and a list of words ("names")

    # Process the rankings and the PPMI line-by-line so as to not run out of memory
    ndcg_a = []
    #spearman_a = []
    with open(rankings_fn) as rankings:
        r = 0
        for lr in rankings:
                for lp in ppmi:
                    # Get the plain-number ranking of the rankings, e.g. "1, 4, 3, 50"
                    sorted_indices = np.argsort(list(map(float, lr.strip().split())))[::-1]
                    # Convert PPMI scores to floats
                    # Get the NDCG score for the PPMI score, which is a valuation, compared to the indice of the rank
                    ndcg = ndcg_from_ranking(lp, sorted_indices)

                    # Add to array and print
                    ndcg_a.append(ndcg)
                    print("ndcg", ndcg, names[r], r)
                    """
                    smr = spearmanr(ppmi_indices, sorted_indices)[1]
                    spearman_a.append(smr)
                    print("spearman", smr, names[r], r)
                    """
                    r+=1
                    break
    # Save NDCG
    dt.write1dArray(ndcg_a, ndcg_fn)
    #dt.write1dArray(spearman_a, spearman_fn)



class GetNDCG(Method.Method):
    word_dir = None
    words = None
    output_directions = None
    output_folder = None
    directions = None
    ranks = None
    bowmin = None
    bowmax = None
    new_word_dict = None
    ppmi = None
    LR = None
    ppmi_id_dct = None

    def __init__(self, ranks, ppmi, words, ppmi_id_dct, save_class, file_name, output_folder, bowmin, bowmax):
        self.words = words
        self.output_folder = output_folder
        self.ppmi = ppmi
        self.ranks = ranks
        self.bowmin = bowmin
        self.bowmax = bowmax
        self.ppmi_id_dct = ppmi_id_dct
        super().__init__(file_name, save_class)

    def makePopos(self):
        self.ndcg_scores = SaveLoadPOPO(np.empty(len(self.ranks), dtype=np.float64),
                                        self.output_folder  + self.file_name + "_" + str(
                                            self.bowmin) + "_" + str(self.bowmax) + "_ndcg.npy", "npy")
        # If the rankings for this situation exists, then dont load the overall.
        self.ndcg_dir = SaveLoadPOPO({}, self.output_folder + self.file_name + "_ndcg_dir.npy", "npy_dict")
        if self.save_class.exists([self.ndcg_dir]) and self.save_class.exists([self.ndcg_scores]) is False:
            self.ndcg_dir.value = self.save_class.load(self.ndcg_dir)
            print("pred_dir loaded successfully")
        self.csv_data = SaveLoadPOPO([["ndcg"],[],[list(self.words.keys())]],
                                        self.output_folder + "csv/"+ self.file_name + "_" + str(
                                            self.bowmin) + "_" + str(self.bowmax) + "_ndcg.csv","csv")


    def makePopoArray(self):
        self.popo_array = [self.ndcg_scores, self.ndcg_dir, self.csv_data]

    def process(self):
        """
        for i in range(len(self.words)):
            if self.words[i] not in self.ndcg_dir.value:
                sorted_indices = np.flipud(np.argsort(self.ranks[i]))
                # Get the NDCG score for the PPMI score, which is a valuation, compared to the indice of the rank
                ndcg = ndcg_from_ranking(np.asarray(self.ppmi[self.ppmi_id_dct[self.words[i]]].todense())[0], sorted_indices)
                self.ndcg_dir.value[self.words[i]] = ndcg
            self.ndcg_scores.value[i] = self.ndcg_dir.value[self.words[i]]
            print(i, "/", len(self.words), self.words[i], self.ndcg_dir.value[self.words[i]])#
        """
        i = 0
        for key, value in self.words.items():
            if key not in self.ndcg_dir.value:
                sorted_indices = np.flipud(np.argsort(self.ranks[value]))
                # Get the NDCG score for the PPMI score, which is a valuation, compared to the indice of the rank
                ndcg = ndcg_from_ranking(np.asarray(self.ppmi[self.ppmi_id_dct[key]].todense())[0],
                                         sorted_indices)
                self.ndcg_dir.value[key] = ndcg
            self.ndcg_scores.value[value] = self.ndcg_dir.value[key]
            print(i, "/", len(self.words.keys()), key, self.ndcg_dir.value[key])
            i += 1

        self.csv_data.value[1] = [self.ndcg_scores.value]
        super().process()

    def getNDCG(self):
        if self.processed is False:
            self.ndcg_scores.value = self.save_class.load(self.ndcg_scores)
        return self.ndcg_scores.value

class GetNDCGStreamed(Method.Method):
    word_dir = None
    words = None
    output_directions = None
    output_folder = None
    directions = None
    ranks_fn = None
    bowmin = None
    bowmax = None
    new_word_dict = None
    ppmi = None
    LR = None
    ppmi_id_dct = None

    def __init__(self, ranks_fn, ppmi, words, ppmi_id_dct, save_class, file_name, output_folder, bowmin, bowmax):
        self.words = words
        self.output_folder = output_folder
        self.ppmi = ppmi
        self.ranks_fn = ranks_fn
        self.bowmin = bowmin
        self.bowmax = bowmax
        self.ppmi_id_dct = ppmi_id_dct
        super().__init__(file_name, save_class)

    def makePopos(self):
        self.ndcg_scores = SaveLoadPOPO(np.empty(len(list(self.words.keys())), dtype=np.float64),
                                        self.output_folder  + self.file_name + "_" + str(
                                            self.bowmin) + "_" + str(self.bowmax) + "Streamed_ndcg.npy", "npy")
        # If the rankings for this situation exists, then dont load the overall.
        self.ndcg_dir = SaveLoadPOPO({}, self.output_folder + self.file_name + "Streamed_ndcg_dir.npy", "npy_dict")
        if self.save_class.exists([self.ndcg_dir]) and self.save_class.exists([self.ndcg_scores]) is False:
            self.ndcg_dir.value = self.save_class.load(self.ndcg_dir)
            print("pred_dir loaded successfully")
        self.csv_data = SaveLoadPOPO([["ndcg"],[],[list(self.words.keys())]],
                                        self.output_folder + "csv/"+ self.file_name + "_" + str(
                                            self.bowmin) + "_" + str(self.bowmax) + "Streamed_ndcg.csv","csv")


    def makePopoArray(self):
        self.popo_array = [self.ndcg_scores, self.ndcg_dir, self.csv_data]


    def process(self):
        j = 0
        keys = np.asarray(list(self.words.keys()))
        vals = np.asarray(list(self.words.values()))
        sorted_ids = np.argsort(vals)
        sorted_keys = keys[sorted_ids]
        sorted_vals = vals[sorted_ids]
        print(self.ranks_fn)
        with open(self.ranks_fn) as infile:
            i = 0
            prev_sorted_val = 0
            for line in infile:
                if len(line) < 1:
                    continue
                if sorted_keys[i] not in self.ndcg_dir.value:
                    if (sorted_vals[i] - prev_sorted_val) > 1:
                        raise ValueError("Dict has missing value, cannot write in rows")
                    rank_str = line.split()
                    rank_float = np.empty(len(rank_str), dtype=np.float64)
                    for j in range(len(rank_str)):
                        rank_float[j] = float(rank_str[j])

                    sorted_indices = np.flipud(np.argsort(rank_float))
                    # Get the NDCG score for the PPMI score, which is a valuation, compared to the indice of the rank
                    try:
                        ndcg = ndcg_from_ranking(np.asarray(self.ppmi[self.ppmi_id_dct[sorted_keys[i]]].todense())[0],
                                             sorted_indices)
                    except IndexError:
                        ndcg = 0.0
                        print(sorted_indices)
                        print(self.ppmi[self.ppmi_id_dct[sorted_keys[i]]].todense()[0])
                    self.ndcg_dir.value[sorted_keys[i]] = ndcg
                self.ndcg_scores.value[sorted_vals[i]] = self.ndcg_dir.value[sorted_keys[i]]
                print(i, "/", len(self.words.keys()), sorted_keys[i], self.ndcg_dir.value[sorted_keys[i]])
                prev_sorted_val = sorted_vals[i]
                i += 1
        self.csv_data.value[1] = [self.ndcg_scores.value]
        super().process()

    def getNDCG(self):
        if self.processed is False:
            self.ndcg_scores.value = self.save_class.load(self.ndcg_scores)
        return self.ndcg_scores.value

class GetNDCGStreamedSimple(Method.Method):
    word_dir = None
    words = None
    output_directions = None
    output_folder = None
    directions = None
    ranks_fn = None
    bowmin = None
    bowmax = None
    new_word_dict = None
    ppmi = None
    LR = None
    ppmi_id_dct = None

    def __init__(self, ranks_fn, ppmi,  words, save_class, file_name, output_folder):
        self.output_folder = output_folder
        self.ppmi = ppmi
        self.ranks_fn = ranks_fn
        self.words=words
        super().__init__(file_name, save_class)

    def makePopos(self):
        try:
            self.ndcg_scores = SaveLoadPOPO(np.empty(len(list(self.words.keys())), dtype=np.float64),
                                        self.output_folder  + self.file_name + "_" + str(
                                            self.bowmin) + "_" + str(self.bowmax) + "Streamed_ndcg.npy", "npy")
        except:
            self.ndcg_scores = SaveLoadPOPO(np.empty(len(self.words), dtype=np.float64),
                                            self.output_folder + self.file_name + "_" + str(
                                                self.bowmin) + "_" + str(self.bowmax) + "Streamed_ndcg.npy", "npy")
        # If the rankings for this situation exists, then dont load the overall.
        self.ndcg_dir = SaveLoadPOPO({}, self.output_folder + self.file_name + "Streamed_ndcg_dir.npy", "npy_dict")
        if self.save_class.exists([self.ndcg_dir]) and self.save_class.exists([self.ndcg_scores]) is False:
            self.ndcg_dir.value = self.save_class.load(self.ndcg_dir)
            print("pred_dir loaded successfully")
        try:
            self.csv_data = SaveLoadPOPO([["ndcg"],[],[list(self.words.keys())]],
                                        self.output_folder + "csv/"+ self.file_name + "_" + str(
                                            self.bowmin) + "_" + str(self.bowmax) + "Streamed_ndcg.csv","csv")
        except:
            self.csv_data = SaveLoadPOPO([["ndcg"],[],[self.words]],
                                        self.output_folder + "csv/"+ self.file_name + "_" + str(
                                            self.bowmin) + "_" + str(self.bowmax) + "Streamed_ndcg.csv","csv")



    def makePopoArray(self):
        self.popo_array = [self.ndcg_scores, self.ndcg_dir, self.csv_data]


    def process(self):
        j = 0

        print(self.ranks_fn)
        with open(self.ranks_fn) as infile:
            i = 0
            for line in infile:
                rank_str = line.split()
                rank_float = np.empty(len(rank_str), dtype=np.float64)
                for j in range(len(rank_str)):
                    rank_float[j] = float(rank_str[j])

                sorted_indices = np.flipud(np.argsort(rank_float))
                # Get the NDCG score for the PPMI score, which is a valuation, compared to the indice of the rank
                ndcg = ndcg_from_ranking(np.asarray(self.ppmi[self.ppmi_id_dct[i]].todense())[0],  sorted_indices)
                self.ndcg_dir.value[i] = ndcg
                self.ndcg_scores.value[i] = self.ndcg_dir.value[i]
                print(i, "/", len(self.words.keys()), i, self.ndcg_dir.value[i])

                i += 1
        self.csv_data.value[1] = [self.ndcg_scores.value]
        super().process()

    def getNDCG(self):
        if self.processed is False:
            self.ndcg_scores.value = self.save_class.load(self.ndcg_scores)
        return self.ndcg_scores.value



if __name__ == '__main__':
    print("kay")