# vim: ft=python

from Compiler import types

from Compiler.util import *

from Compiler.oram import OptimalORAM

from Compiler.library import for_range, do_while, time, if_, print_ln, crash, print_str, break_loop
from Compiler.gs import OMatrix, OStack


class Matchmaker:
    """
    Based on Matchmaker from Compiler/gs.py in MP-SPDZ, copyright (c) 2023,
    Commonwealth Scientific and Industrial Research Organisation (CSIRO)
    ABN 41 687 119 230, published under the BSD 3-Clause Licence
    """

    def pair(self, patient, therapist, for_real):
        self.paired_therapists.access(patient, therapist, for_real)
        self.paired_patients.access(therapist, patient, for_real)

    def unpair(self, patient, therapist, for_real):
        self.paired_therapists.delete(patient, for_real)
        self.paired_patients.delete(therapist, for_real)
        self.unpaired.append(patient, for_real)

    def request_therapist(self, patient, therapist, for_real):
        experience = self.t_exps[therapist][0]
        (old_patient,), free = self.paired_patients.read(therapist)  # patient paired to therapist
        paired = 1 - free
        rank_patient = self.p_cases[patient][0]
        rank_old_patient = self.p_cases[old_patient][0] * paired
        matches_exp = self.int_type(rank_patient) == self.int_type(experience)
        same_as_old = self.int_type(rank_patient) != self.int_type(rank_old_patient)
        leaving = paired * matches_exp * same_as_old

        print_str('therapist: %s, patient: %s, old patient: %s, ',
                  *(x.reveal() for x in (therapist, patient, old_patient)))
        print_ln('rank patient: %s, rank old patient: %s, paired: %s, leaving: %s',
                 *(x.reveal() for x in
                   (rank_patient, rank_old_patient, paired, leaving)))

        self.unpair(old_patient, therapist, paired * leaving * for_real)
        self.pair(patient, therapist, (1 - (paired * (1 - leaving))) * for_real)
        self.unpaired.append(patient, paired * (1 - leaving) * for_real)

    def match(self, n_loops=None):
        if n_loops is None or n_loops > self.N * self.M:
            loop = do_while
            init_rounds = self.N
        else:
            loop = for_range(n_loops)
            init_rounds = n_loops / self.M

        self.paired_therapists = \
            self.oram_type(self.N, entry_size=log2(self.N),
                           init_rounds=0, value_type=self.basic_type)
        self.paired_patients = \
            self.oram_type(self.N, entry_size=log2(self.N),
                           init_rounds=0, value_type=self.basic_type)
        self.unpaired = OStack(self.N, oram_type=self.oram_type,
                           int_type=self.int_type)

        @for_range(init_rounds)
        def _(i):
            self.unpaired.append(i)
        rounds = types.MemValue(types.regint(0))

        @loop
        def _(i=None):
            rounds.iadd(1)
            time()
            patient = self.unpaired.pop()
            pref = self.int_type(p_cases[patient][0])  # patient suffers from pref
            therapist = types.MemValue(self.int_type(0))

            # Get index of first free therapist who has experience with pref
            @for_range(self.N)
            def _(i):
                (_,), free = self.paired_patients.read(i)
                @if_(((pref == t_exps[i][0]) * free).reveal())
                def _():
                    therapist.write(self.int_type(i))
                    break_loop()

                @if_((i == self.N).reveal())
                def _():
                    print_ln('run out of acceptable therapists')
                    crash()

            self.request_therapist(patient, therapist.read(), True)
            print_ln('patient: %s, pref: %s, left: %s',
                     *(x.reveal() for x in (patient, pref, self.unpaired.size)))
            return types.regint((self.unpaired.size > 0).reveal())

        print_ln('%s rounds', rounds)
        print_ln('\nPRINTING PAIRS\n')

        @for_range(init_rounds)
        def _(i):
            print_ln('patient %s : therapist %s', i, self.paired_therapists[i].reveal())

    def __init__(self, p_cases, t_exps, N, M=None, reverse=False,
                 oram_type=OptimalORAM, int_type=types.sint):
        self.N = N
        self.M = N if M is None else M
        self.p_cases = p_cases
        self.t_exps = t_exps
        self.reverse = reverse
        self.oram_type = oram_type
        self.int_type = int_type
        self.basic_type = int_type.basic_type


# Constants
PLAYERS = 3
MATCHING_SIZE = 50

p_shares = Matrix(rows=PLAYERS, columns=MATCHING_SIZE, value_type=types.sint)
t_shares = Matrix(rows=PLAYERS, columns=MATCHING_SIZE, value_type=types.sint)

# Fill data from players into the matrices
@for_range(PLAYERS)
def _(i):
    @for_range(2 * MATCHING_SIZE)
    def _(j):
        index = j % MATCHING_SIZE
        typ = sint.get_input_from(i)

        @if_((typ == -100).reveal())
        def _():
            p_shares[i][index] = sint.get_input_from(i)

        @if_((typ == -200).reveal())
        def _():
            t_shares[i][index] = sint.get_input_from(i)

# Add entire columns together to recreate secret-shared input in ORAM
p_cases = OMatrix(N=MATCHING_SIZE, M=1, oram_type=OptimalORAM, int_type=types.sint)
t_exps = OMatrix(N=MATCHING_SIZE, M=1, oram_type=OptimalORAM, int_type=types.sint)

@for_range(MATCHING_SIZE)
def _(i):
    p_cases[i][0] = sum(p_shares.get_column(i))
    print_ln('patient %s suffers from case %s', i, p_cases[i][0].reveal())
    t_exps[i][0] = sum(t_shares.get_column(i))
    print_ln('therapist %s is experienced with case %s', i, t_exps[i][0].reveal())

# Run algorithm
mm = Matchmaker(p_cases, t_exps, N=MATCHING_SIZE, M=1)
mm.match()