import numpy as np
from numba import jit, njit
from itertools import permutations
from tools import combinations_of_2
from tqdm import tqdm
from tprint import tprint

Number = int|float

# original, unoptimized version, but it's more readable
# def kendall_tau_dist(rank_a, rank_b) -> int:
#     tau = 0
#     nb_candidates = len(rank_a)
#     # for each pair of candidates (not counting twice pairs with a different order)
#     for i, j in combinations(range(nb_candidates), 2):
#         tau += np.sign(rank_a[i] - rank_a[j]) == - np.sign(rank_b[i] - rank_b[j])
#     return tau




def kendall_tau_dist(ranking_a: list[int], ranking_b: list[int]) -> Number:
    """The kendall τ distance between two rankings / permutations.
    It is the number of inversions that don't have the same sign within all pairs of an inversion of ranking_a and an inversion of ranking_b.
    """
    ranking_a = np.array(ranking_a)
    ranking_b = np.array(ranking_b)
    # comb = np.array([[i, j] for i in range(1, size) for j in range(1, size) if i<j])
    comb = combinations_of_2(ranking_a.size)
    A = ranking_a[comb]
    B = ranking_b[comb]

    length = ranking_a.size
    # return __tau(A, B)
    return 1 - __tau(A, B) / (length * (length - 1))


@jit(cache=True)
def __tau(A: list[int], B: list[int]) -> int:
    # tau = np.sum(np.sign(A[:,0] - A[:,1]) == - np.sign(B[:,0] - B[:,1]))
    return np.sum((A[:,0] < A[:,1]) == (B[:,0] > B[:,1]))


def rank_aggregation(rankings: list[list[int]]) -> tuple[int, tuple[int, ...]]:
    """Brute-force kemeny-young rank aggregation.
    Args:
        ranks: A list of the ranks (2D numpy array).
    Returns:
        int, list: The minimal sum of distances to ranks, the rank of minimal distance.
    """
    rankings = np.array(rankings)
    min_dist: int = np.inf
    best_ranking = None
    _, nb_candidates = rankings.shape

    for candidate_ranking in permutations(range(nb_candidates)):
        # dist = sum(kendall_tau_dist(candidate_ranking, ranking) for ranking in rankings)
        dist = np.sum(np.apply_along_axis(lambda rkng: kendall_tau_dist(candidate_ranking, rkng), 1, rankings))
        # dist = np.sum(np.apply_along_axis(lambda rank: kendalltau_dist(candidate_rank, rank), 1, ranks))
        if dist < min_dist:
            min_dist = dist
            best_ranking = candidate_ranking

    # if no minimum has been found, there is a problem somewhere
    if best_ranking is None: raise ValueError()

    return min_dist, best_ranking


if __name__ == '__main__':
    ranks = np.array([[0, 1, 2, 3, 4],
                      [0, 1, 3, 2, 4],
                      [4, 1, 2, 0, 3],
                      [4, 1, 0, 2, 3],
                      [4, 1, 3, 2, 0]])

    # print(rank_aggregation(ranks))

    # print(kendall_tau_dist([1, 2, 3],
    #                        [3, 1, 2]))

    ranks = np.array(list(permutations(range(7))))
    for _ in tqdm(range(10)):
        selected_lines = np.random.randint(ranks.shape[0], size=30)
        selected = ranks[selected_lines,:]
        print(rank_aggregation(selected))
        # tprint(selected)
        # print(ranks)
        # print(kendalltau_dist(ranks[5], ranks[-1]))
        # print(np_kendalltau_dist(ranks[5], ranks[-1]))


