from pathfinder.core.model.Node import Node
from pathfinder.core.repo.Repository import Repository


class NGDSortedNeighborsRepo(Repository):

    def __init__(self, repo, degree_repo, ngd_repo):
        self.repo = repo
        self.degree_repo = degree_repo
        self.ngd_repo = ngd_repo

    def get_neighbors(self, node, limit=-1):
        if limit <= 0:
            raise Exception(f"The limit:{limit} could not be negative or zero.")
        curie_ngd_list = self.ngd_repo.get_curie_ngd(node.id)
        ngd_by_curie_dict = {}
        if curie_ngd_list is not None:
            if len(curie_ngd_list) >= limit:
                return [Node(curie, ngd_value) for curie, ngd_value in
                        curie_ngd_list[0:min(limit, len(curie_ngd_list))]]
            else:
                for curie, ngd in curie_ngd_list:
                    ngd_by_curie_dict[curie] = ngd

        id_set = set(ngd_by_curie_dict.keys())
        neighbors = self.repo.get_neighbors(node, limit=limit)
        neighbors_ids = []
        for neighbor in neighbors:
            if neighbor.id not in id_set:
                neighbors_ids.append(neighbor.id)
            else:
                id_set.remove(neighbor.id)

        for remove_id in id_set:
            del ngd_by_curie_dict[remove_id]

        number_of_curie_left_to_fill_the_limit = limit - len(curie_ngd_list)

        node_pmids_length = self.ngd_repo.get_curies_pmid_length(neighbors_ids,
                                                                 number_of_curie_left_to_fill_the_limit)
        node_pmids_length = [(key, length) for key, length in node_pmids_length if length > 0]

        curies_sorted_by_their_pmids_length = sorted(node_pmids_length, key=lambda x: (x[1]), reverse=True)

        nonzero_length_pmids_curie_with_none_ngd_value = [(curie, None) for curie, _ in
                                                          curies_sorted_by_their_pmids_length]

        sorted_neighbors_tuple = sorted(ngd_by_curie_dict.items(),
                                        key=lambda x: (x[1] is None, x[1] if x[1] is not None else float('inf')))
        sorted_neighbors_tuple.extend(nonzero_length_pmids_curie_with_none_ngd_value)

        return [Node(item[0], item[1]) for item in
                sorted_neighbors_tuple[0:min(limit, len(sorted_neighbors_tuple))]]

    def get_node_degree(self, node_id):
        return self.degree_repo.get_node_degree(node_id)
