import numpy as np
import numba as nb
from itertools import permutations
from fastcache import lru_cache
from tqdm import tqdm
from tprint import tprint

@nb.jit
def simple_borda(rankings: list[list[int]]) -> list[int]:
    # dictionnary to count the points given to each value
    points_count = {value: 0 for value in rankings[0]}
    length = len(rankings[0])
    for ranking in rankings:
        for points, value in enumerate(ranking):
            points_count[value] += length - points
    # sort dictionnary items by values
    sorted_points_count = sorted(points_count.items(),
                                 key=lambda x: x[1],
                                 reverse=True)
    # return dictionnary keys
    return [item[0] for item in sorted_points_count]

@nb.jit
def borda(rankings: list[list[int]]) -> list[int]:
    length = len(rankings[0])
    points_count = np.zeros(length)
    for ranking in rankings:
        points_count += np.argsort(ranking)
    return np.argsort(points_count)





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(borda(ranks))

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

    def take(idx: int, iterator: iter) -> iter:
        for _ in range(idx):
            yield next(iterator)

    ranks = np.array(list(take(20000, permutations(range(20)))))
    for _ in tqdm(range(500)):
        selected_lines = np.random.randint(ranks.shape[0], size=20000)
        selected = ranks[selected_lines,:]
        if (borda(selected) != simple_borda(selected)).any():
            print(*borda(selected))
            print(*simple_borda(selected))
        # tprint(selected)
        # print(ranks)
        # print(kendalltau_dist(ranks[5], ranks[-1]))
        # print(np_kendalltau_dist(ranks[5], ranks[-1]))


