from dataset.AbstractDataset import AbstractDataset
import numpy as np
import math
import os
import hashlib
import pickle


class LetorDataset(AbstractDataset):

    def __init__(self,
                 path,
                 feature_size,
                 query_level_norm=False,
                 binary_label=0,
                 cache_root=None
                 ):
        super().__init__(path, feature_size, query_level_norm)
        self._binary_label = binary_label
        self._comments = {}
        self._docid_map = {}
        # self._docstr_map = {}

        if cache_root is not None:
            new = self.from_path(path, cache_root)
            if new is not None:
                self.__dict__.update(new.__dict__)
            else:
                self._load_data()
                cache_name = f"{hashlib.md5(path.encode()).hexdigest()}.pkl"
                file_path = f'./{cache_root}/{cache_name}'

                with open(file_path, 'wb') as f:
                    pickle.dump(self, f)
                print(f"Cached the array to {file_path}")
        else:
            self._load_data()

    def from_path(self, path: str, cache_root: str) -> 'LetorDataset':
        """
        Constructs a dataset by reading form a disk
        :param root_path: A path to the root that contains (Fold1, Fold2, ...)
        :param cache_root: None if no caching is needed, otherwise a path to the cache dir;
                        if the cache dir already contains the data it will be used. Hence, cleaning
                        of the cache has to be done manually if needed.
        :return: Constructed Dataset instance
        """

        cache_name = f"{hashlib.md5(path.encode()).hexdigest()}.pkl"
        file_path = f'./{cache_root}/{cache_name}'
        if os.path.exists(file_path):
            print(f"Loading from cache file {file_path}")
            with open(file_path, 'rb') as f:
                data = pickle.load(f)
            return data
        else:
            print("no cache found for", path)
            return None

    def _load_data(self):
        print("Loading {}......".format(self._path))
        with open(self._path, "r") as fin:
            current_query = None
            for line in fin:
                cols = line.strip().split()
                query = cols[1].split(':')[1]
                if query == current_query:
                    docid = len(self._query_get_docids[query])
                    old_query = True

                else:
                    if current_query != None and self._query_level_norm:
                        self._normalise(current_query)
                    old_query = False
                    docid = 0
                    current_query = query
                    self._docid_map[query] = {}
                    # self._docstr_map[query] = {}
                    self._query_pos_docids[query] = []

                comments_part = line.split("#")
                if len(comments_part) == 2:
                    if query not in self._comments:
                        self._comments[query] = []
                    self._comments[query].append(comments_part[1].strip())

                relevence = float(cols[0])  # Sometimes the relevance label can be a float.
                if relevence.is_integer():
                    relevence = int(relevence)  # But if it is indeed an int, cast it into one.
                if self._binary_label != 0:
                    if relevence >= self._binary_label:
                        relevence = 1
                    else:
                        relevence = 0

                features = [0] * self._feature_size

                for i in range(2, len(cols)):
                    feature_id = cols[i].split(':')[0]

                    if not feature_id.isdigit():
                        if feature_id[0] == "#":
                            self._docid_map[query][docid] = cols[i][1:]
                            # self._docstr_map[query][cols[i][1:]] = docid
                        break

                    feature_id = int(feature_id) - 1
                    feature_value = float(cols[i].split(':')[1])
                    if math.isnan(feature_value):
                        feature_value = 0

                    features[feature_id] = feature_value

                if relevence > 0:
                    self._query_pos_docids[query].append(docid)

                if old_query:
                    self._query_docid_get_features[query][docid] = np.array(features)
                    self._query_get_docids[query].append(docid)
                    self._query_get_all_features[query] = np.vstack((self._query_get_all_features[query], features))
                    self._query_docid_get_rel[query][docid] = relevence
                    self._query_relevant_labels[query].append(relevence)
                else:
                    self._query_docid_get_features[query] = {docid: np.array(features)}
                    self._query_get_docids[query] = [docid]
                    self._query_get_all_features[query] = np.array([features])
                    self._query_docid_get_rel[query] = {docid: relevence}
                    self._query_relevant_labels[query] = [relevence]

        if self._query_level_norm:
            self._normalise(current_query)

    def _normalise(self, query):
        norm = np.zeros((len(self._query_get_docids[query]), self._feature_size))
        # if there is more than 1 candidate docs, do the norm
        if norm.shape[0] != 1:
            query_features = self._query_get_all_features[query]
            min = np.amin(query_features, axis=0)
            max = np.amax(query_features, axis=0)
            safe_ind = max - min != 0
            norm[:, safe_ind] = (query_features[:, safe_ind] - min[safe_ind]) / (
                    max[safe_ind] - min[safe_ind])
            self._query_get_all_features[query] = norm

    def update_relevance_label(self, qrel_dic: dict):
        for qid in self._query_docid_get_rel.keys():

            self._query_pos_docids[qid] = []
            ind = 0
            for docid in self._query_docid_get_rel[qid].keys():
                if self._docid_map[qid][docid] in qrel_dic[qid].keys():
                    rel = qrel_dic[qid][self._docid_map[qid][docid]]
                else:
                    rel = 0
                self._query_docid_get_rel[qid][docid] = rel
                self._query_relevant_labels[qid][ind] = rel

                if rel > 0:
                    self._query_pos_docids[qid].append(docid)
                ind += 1

        # self._query_get_all_features_temp = {}
        # for qid in self._query_docid_get_rel.keys():
        #     self._query_pos_docids[qid] = []
        #     self._query_get_docids[qid] = []
        #     self._query_get_all_features_temp[qid] = []
        #     for docstr in qrel_dic[qid].keys():
        #         docid = self._docstr_map[qid][docstr]
        #         rel = qrel_dic[qid][docstr]
        #
        #         self._query_docid_get_rel[qid][docid] = rel
        #         self._query_relevant_labels[qid][docid] = rel
        #         self._query_get_docids[qid].append(docid)
        #         features = self._query_get_all_features[qid][docid]
        #         self._query_get_all_features_temp[qid].append(features)
        #         if rel > 0:
        #             self._query_pos_docids[qid].append(docid)
        #     self._query_get_all_features[qid] = np.array(self._query_get_all_features_temp[qid])

    def update_relevance_by_qrel(self, path: str):

        # q-d pair dictionary
        qrel_dic = {}

        with open(path, 'r') as f:
            for line in f:
                qid, _, docid, rel = line.strip().split()
                if qid in qrel_dic.keys():
                    qrel_dic[qid][docid] = int(rel)
                else:
                    qrel_dic[qid] = {docid: int(rel)}
        self.update_relevance_label(qrel_dic)

    def get_features_by_query_and_docid(self, query, docid):
        return self._query_docid_get_features[query][docid]

    def get_candidate_docids_by_query(self, query):
        return self._query_get_docids[query]

    def get_all_features_by_query(self, query):
        return self._query_get_all_features[query]

    def get_relevance_label_by_query_and_docid(self, query, docid):
        return self._query_docid_get_rel[query][docid]

    def get_all_relevance_label_by_query(self, query):
        return self._query_relevant_labels[query]

    def get_relevance_docids_by_query(self, query):
        return self._query_pos_docids[query]

    def get_all_querys(self):
        return np.array(list(self._query_get_all_features.keys()))

    def get_all_comments_by_query(self, query):
        return self._comments[query]

    def write(self, output_file):
        s = ""
        for query in self.get_all_querys():
            comments = self.get_all_comments_by_query(query)
            for i, features in enumerate(self.get_all_features_by_query(query)):
                comment = comments[i]
                label = self.get_relevance_label_by_query_and_docid(query, i)
                features_str = ""
                for i, feature in enumerate(features):
                    # if feature == 0:
                    #     continue
                    features_str += "{}:{} ".format(i + 1, feature)
                s += "{} qid:{} {}#{}\n".format(label, query, features_str, comment)
        with open(output_file, "w") as f:
            f.write(s)

    def write_by_queries(self, output_file, queries):
        s = ""
        for query in queries:
            comments = self.get_all_comments_by_query(query)
            for i, features in enumerate(self.get_all_features_by_query(query)):
                comment = comments[i]
                label = self.get_relevance_label_by_query_and_docid(query, i)
                features_str = ""
                for i, feature in enumerate(features):
                    # if feature == 0:
                    #     continue
                    features_str += "{}:{} ".format(i + 1, feature)
                s += "{} qid:{} {}#{}\n".format(label, query, features_str, comment)
        with open(output_file, "w") as f:
            f.write(s)

    def write_cross_validation_datasets(self, path: str, fold_num: int):
        """
        :param fold_num: number of fold to do cross validation.
        :param path: folder address to store the cross sets.
        :return:
        """

        for fold in range(fold_num):
            fold_path = "{}/fold{}".format(path, fold + 1)
            # Create target Directory if don't exist
            if not os.path.exists(fold_path):
                os.mkdir(fold_path)
                print("Directory ", fold_path, " Created ")
            else:
                print("Directory ", fold_path, " already exists")

        all_queries = self.get_all_querys()

        np.random.shuffle(all_queries)

        query_chunks = np.array_split(all_queries, fold_num)

        for i in range(fold_num):
            test_path = "{}/Fold{}/test.txt".format(path, i + 1)
            train_path = "{}/Fold{}/train.txt".format(path, i + 1)
            test_queries = query_chunks[i]
            train_chunks = query_chunks[:i]
            train_chunks.extend(query_chunks[i + 1:])
            train_queries = np.concatenate(train_chunks)

            self.write_by_queries(test_path, test_queries)
            self.write_by_queries(train_path, train_queries)

    @staticmethod
    def runs_to_letor(input_folder: str, output_folder: str):
        """
        Convert run files into LTR dataset.
        :param input_folder: folder path that contains all run files.
        :param output_folder:
        :return:
        """
        files = os.listdir(input_folder)
        num_feature = len(files)

        # q-d pair dictionary
        query_dic = {}
        for feature_id in range(num_feature):
            # feature id in standard letor datasets start from 1.
            with open(os.path.join(input_folder, files[feature_id]), 'r') as f:
                for line in f:
                    qid, _, docid, rank, score, rname = line.strip().split()
                    if qid in query_dic.keys():
                        if docid in query_dic[qid].keys():
                            query_dic[qid][docid].append((feature_id + 1, score))
                        else:
                            query_dic[qid][docid] = [(feature_id + 1, score)]
                    else:
                        query_dic[qid] = {docid: [(feature_id + 1, score)]}
        s = ""
        for qid in query_dic.keys():
            for docid in query_dic[qid].keys():
                # the first column is relevance label, dont know for now.
                s += "0 "
                s += "qid:{} ".format(qid)

                for feature_id, socre in query_dic[qid][docid]:
                    s += "{}:{} ".format(feature_id, socre)
                s += "#docid = {}\n".format(docid)
        with open(output_folder + "letor.txt", "w") as f:
            f.write(s)

        s = ""
        for fid in range(len(files)):
            s += "{}:{}\n".format(fid + 1, files[fid])
        with open(output_folder + "feature_description.txt", "w") as f:
            f.write(s)

    # bad implementation only for PMGD:
    def get_query_docid_get_feature(self):
        return self._query_docid_get_features

    def get_query_get_all_features(self):
        return self._query_get_all_features

    def get_query_get_docids(self):
        return self._query_get_docids
