#!/usr/bin/env python3
import argparse
import json
import numpy as np

# N[a,b] is the number of voters who prefer candidate a to candidate b
def load_pairwise_preferences_matrix(file):
    list_of_lists = json.load(file)
    n = np.stack([np.asarray(l, dtype=np.int8) for l in list_of_lists])
    return n

# As defined in section 4.12.1 https://arxiv.org/pdf/1804.02973.pdf,
# A Condorcet winner is an alternative a ∈ A that wins every head-to-head
# contest with some other alternative b ∈ A \ {a}.
def condorcet_winner(n, candidates):
    nb_candidates = len(candidates)
    for potential_winner in range(nb_candidates):
        won_all_contests = True
        for opponent in range(nb_candidates):
            if n[opponent][potential_winner] > n[potential_winner][opponent]:
                won_all_contests = False
        if won_all_contests:
            return candidates[potential_winner]
    return None

# As defined in section 4.12.1 https://arxiv.org/pdf/1804.02973.pdf,
# A weak Condorcet winner is an alternative a ∈ A that doesn’t lose any
# head-to-head contest with some other alternative b ∈ A \ {a}
def weak_condorcet_winners(n, candidates):
    weak_winners = []
    nb_candidates = len(candidates)
    for potential_winner in range(nb_candidates):
        lost_a_contest = False
        for opponent in range(nb_candidates):
            if n[opponent][potential_winner] > n[potential_winner][opponent]:
                lost_a_contest = True
        if not lost_a_contest:
            weak_winners.append(candidates[potential_winner])
    return weak_winners

# As defined in section 4 https://citizensassembly.arts.ubc.ca/resources/submissions/csharman-10_0409201706-143.pdf
# + other link strength methods defined in https://arxiv.org/pdf/1804.02973.pdf
def schulze_winners(n, candidates, link_strength_method='margin'):
    nb_candidates = len(candidates)

    if link_strength_method == 'margin':
        p = np.zeros((nb_candidates, nb_candidates), dtype=np.int8)
        for i in range(nb_candidates):
            for j in range(nb_candidates):
                if i != j:
                    p[i][j] = n[i][j] - n[j][i]
    elif link_strength_method == 'ratio':
        p = np.ones((nb_candidates, nb_candidates), dtype=np.int8) * np.Inf
        for i in range(nb_candidates):
            for j in range(nb_candidates):
                if i != j and n[j][i] > 0:
                    p[i][j] = n[i][j] / n[j][i]
    else:
        raise ValueError(f"link_strength_method '{link_strength_method}' is not implemented")

    for i in range(nb_candidates):
        for j in range(nb_candidates):
            if i != j:
                for k in range(nb_candidates):
                    if i != k and j != k:
                        s = min(p[j][i], p[i][k])
                        if p[j][k] < s:
                            p[j][k] = s

    winners = []
    for i in range(nb_candidates):
        has_lost_directly_or_indirectly = False
        for j in range(nb_candidates):
            if i != j:
                if p[j][i] > p[i][j]:
                    has_lost_directly_or_indirectly = True
        if not has_lost_directly_or_indirectly:
            winners.append(candidates[i])
    return winners


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Computes winners of the election')
    parser.add_argument('-c', '--candidates-file', help='candidates JSON file', type=argparse.FileType('r', encoding='utf-8'), required=True)
    parser.add_argument('matrix_file', help="pairwise preference matric JSON file", type=argparse.FileType('r', encoding='utf-8'))
    args = parser.parse_args()

    n = load_pairwise_preferences_matrix(args.matrix_file)
    assert(n.shape[0] > 0)
    assert(n.shape[0] == n.shape[1])

    candidates = json.load(args.candidates_file)
    assert(len(candidates) == n.shape[0])

    # Condorcet winner (win every head-to-head contest)
    c_winner = condorcet_winner(n, candidates)
    if c_winner is None:
        print('There is no Condorcet winner')
    else:
        print(f'Condorcet winner: {c_winner}')

    # Condorcet weak winner (do not lose any head-to-head contest)
    weak_c_winners = weak_condorcet_winners(n, candidates)
    if len(weak_c_winners) == 0:
        print('There are no weak Condorcet winners')
    else:
        print(f'Weak Condorcet winners: {weak_c_winners}')

    # Schulze winners (win every head-to-head contest, directly or indirectly)
    schulze_winners = schulze_winners(n, candidates)
    assert(len(schulze_winners) > 0)
    print(f'Schulze winners: {schulze_winners}')