import numpy as np
import random
import scipy.stats as stats
import itertools

class Main:
    def __init__(self, threshold):
        self.threshold = threshold

    def flatten(self, lst):
        return [x for y in lst for x in y]

    def almost_int(self, x):
        return abs(x - round(x))<=0.00001

    def bader_ofer_help(self, fractional_results):
        parties = list(fractional_results) # give consistent order to parties
        fractional_mand = np.array([fractional_results[party] for party in parties])
        tot_mandates = fractional_mand.sum()
        assert (self.almost_int(tot_mandates))
        result = np.floor(fractional_mand)
        leftover_mand = int(tot_mandates - result.sum())
        print (tot_mandates, result.sum(), leftover_mand)
        for i in range(leftover_mand):
            bader_ofer_scores = fractional_mand / (result+1)
            odef_winner = np.argmax(bader_ofer_scores)
            result[odef_winner] += 1
        result = {parties[i] : result[i] for i in range(len(result))}
        return result

    def bader_ofer(self, unnorm_percs, unnorm_odafim):
        tot_percs = sum(unnorm_percs.values())
        percs = { party: unnorm_percs[party] * 120 / tot_percs for party in unnorm_percs }
        odafim = set([ (x,y) for x,y in unnorm_odafim.items() if x in percs and y in percs ])
        no_odafim = set(percs) - set(self.flatten(odafim))
        bader_ofer_step1 = {party : percs[party] for party in no_odafim}
        bader_ofer_step1.update({pair : percs[pair[0]] + percs[pair[1]] for pair in odafim })
        step1_res = self.bader_ofer_help(bader_ofer_step1)
        simu_result = {party : step1_res[party] for party in no_odafim}
        for pair in odafim:
            norm_factor2 = step1_res[pair] / (percs[pair[0]] + percs[pair[1]])
            pair_bader_ofer = {party : percs[party] * norm_factor2 for party in pair}
            simu_result.update(self.bader_ofer_help(pair_bader_ofer))
        return simu_result

    def run_simulation(self, survey, odafim1, odafim2, sigma):
        noise_multiplier = 0.05
        #assert int(noise_multiplier) == noise_multiplier
        #assert set(self.flatten(odafim)) <= set(survey)
        #no_odafim = set(survey) - set(self.flatten(odafim))
        parties = list(survey) # ordered list of parties
        tot_mandates = sum(survey.values())
        print ("total=",tot_mandates)
        #assert int(tot_mandates) == tot_mandates
        total = {party : list() for party in parties}
        alt_gains = {party : list() for party in parties}
        num_simulations = 10000
        for _ in range(num_simulations):
            noisy_survey = {}
            lower = 0.0
            upper = 1.0
            shuffled_parties = [ p for p in survey if p != 'other' ]
            random.shuffle(shuffled_parties)
            while True:
                error = False
                for party in shuffled_parties:
                    party_perc = survey[party]
                    party_std = sigma[party] #np.sqrt(party_perc * (1 - party_perc) / 500) if sigma is None or party not in sigma else sigma[party]
                    low = (lower - party_perc)/party_std
                    upp = (upper - party_perc)/party_std
                    noisy_survey[party] = stats.truncnorm.rvs(low, upp, loc = party_perc, scale = party_std)
                    upper -= noisy_survey[party]
                    if upper < 0:
                        error = True
                        break
                if not error:
                    break
            noisy_survey['other'] = 1 - sum(noisy_survey.values())
            #print ("noisy", survey, noisy_survey)
            assert all([v>=0 for v in noisy_survey.values()])
            #norm_factor1 = tot_mandates / sum(noisy_survey.values())
            #noisy_survey = {party : noisy_survey[party] * norm_factor1 for party in noisy_survey}
            simu_result1 = self.perform_bader_ofer(noisy_survey, odafim1)
            simu_result2 = self.perform_bader_ofer(noisy_survey, odafim2)
##            bader_ofer_step1 = {party : noisy_survey[party] for party in no_odafim}
##            bader_ofer_step1.update({pair : noisy_survey[pair[0]] + noisy_survey[pair[1]] for pair in odafim})
##            step1_res = bader_ofer(bader_ofer_step1)
##            simu_result = {party : step1_res[party] for party in no_odafim}
##            for pair in odafim:
##                norm_factor2 = step1_res[pair] / (noisy_survey[pair[0]] + noisy_survey[pair[1]])
##                pair_bader_ofer = {party : noisy_survey[party] * norm_factor2 for party in pair}
##                simu_result.update(bader_ofer(pair_bader_ofer))
            for party in parties:
                total[party].append(simu_result1[party])
                alt_gains[party].append(simu_result2[party] - simu_result1[party])
        result = {party : sum(total[party]) / num_simulations for party in parties}
        total_gains = {party : sorted(alt_gains[party]) for party in parties}
        return result, total_gains

    def compare_simulations(self, survey, odafim1, odafim2, sigma = None):
        normal_results, alt_gains = self.run_simulation(survey, odafim1, odafim2, sigma)
        return alt_gains

    def perform_bader_ofer(self, percs, surplus_vote_agreements):
        totalpercs = sum(percs.values())
        ahuz_hasima = self.threshold * totalpercs
        passed_votes = 0
        notpassed_votes = 0
        for party in percs.keys():
            if party != 'other' and percs[party] >= ahuz_hasima:
                passed_votes += percs[party]
            else:
                notpassed_votes += percs[party]

        moded = passed_votes / 120
        mands = {}
        for party, perc in percs.items():
            if party != 'other' and perc >= ahuz_hasima:
                mands[party] = int(perc / moded)
            else:
                mands[party] = 0

        #print (totalpercs, ahuz_hasima, mands)
        leftover = 120 -  sum(mands.values())

        jointmands = {}
        for party, mand in mands.items():
            if party in surplus_vote_agreements.keys() or party in surplus_vote_agreements.values():
                for p1, p2 in surplus_vote_agreements.items():
                    if p1 == party or p2 == party:
                        jointparty = '+'.join([p1, p2])
                        if p1 in mands and p2 in mands:
                            jointmands[jointparty] = mands[p1] + mands[p2]
                        elif p1 in mands:
                            jointmands[jointparty] = mands[p1]
                        elif p2 in mands:
                            jointmands[jointparty] = mands[p2]
                        else:
                            jointmands[jointparty] = 0
            else:
                jointmands[party] = mand

        partymodeds = {}
        for party, perc in percs.items():
            if party in surplus_vote_agreements.keys() or party in surplus_vote_agreements.values():
                for p1, p2 in surplus_vote_agreements.items():
                    if p1 == party or p2 == party:
                        jointparty = '+'.join([p1, p2])
                        perc1 = percs[p1] if p1 in percs and percs[p1] >= ahuz_hasima else 0
                        perc2 = percs[p2] if p2 in percs and percs[p2] >= ahuz_hasima else 0
                        partymodeds[jointparty] = (perc1 + perc2) / (jointmands[jointparty] + 1)
            else:
                if party != 'other' and perc >= ahuz_hasima:
                    partymodeds[party] = perc / (mands[party] + 1)
                else:
                    partymodeds[party] = 0

        given = {}
        for i in range(0, leftover):
            maxparty = None
            for party, moded in partymodeds.items():
                if maxparty == None or partymodeds[party] > partymodeds[maxparty]:
                    maxparty = party

            if maxparty in given:
                given[maxparty] += 1
            else:
                given[maxparty] = 1

            jointmands[maxparty] += 1
            if '+' in maxparty:
                parties = maxparty.split('+')
                perc1 = percs[parties[0]] if parties[0] in percs else 0
                perc2 = percs[parties[1]] if parties[1] in percs else 0
                partymodeds[maxparty] = (perc1 + perc2) / (jointmands[maxparty] + 1)
            else:
                partymodeds[maxparty] = percs[maxparty] / (jointmands[maxparty] + 1)

        finalmands = {}
        for party, mand in jointmands.items():
            if '+' in party:
                parties = party.split('+')
                party1 = parties[0]
                party2 = parties[1]
                perc1 = percs[party1] if party1 in percs else 0
                perc2 = percs[party2] if party2 in percs else 0
                if jointmands[party] == 0:
                    finalmands[party1] = 0
                    finalmands[party2] = 0
                else:
                    moded = (perc1 + perc2) / jointmands[party]
                    mands1 = int(perc1 / moded)
                    mands2 = int(perc2 / moded)
                    finalmands[party1] = mands1
                    finalmands[party2] = mands2
                if party in given  and (perc1 >= ahuz_hasima or perc2 >= ahuz_hasima):
                    #print (party, given[party], mands1, mands2)
                    togive = mand - mands1 - mands2
                    for i in range(0, togive):
                        moded1 = perc1 / (mands1 + 1)
                        moded2 = perc2 / (mands2 + 1)
                        if (moded1 > moded2 or finalmands[party2] == 0) and finalmands[party1] > 0:
                            finalmands[party1] += 1
                            mands1 += 1
                        elif finalmands[party2] > 0:
                            finalmands[party2] += 1
                            mands2 += 1
                else:
                    finalmands[party1] = mands[party1] if party1 in mands else 0
                    finalmands[party2] = mands[party2] if party2 in mands else 0
            else:
                finalmands[party] = mand

        if sum(finalmands.values()) != 120:
            print ("percs", percs)
            print ("surplus", surplus_vote_agreements)
            print (sum(finalmands.values()), finalmands)
        assert(sum(finalmands.values()) == 120)
        return finalmands

surplus1 = {'likud': 'bayityehudi', 'kadima': 'amshalem', 'shas': 'yahadut', 'tenuah':'meretz', 'avoda':'yeshatid', 'hadash':'balad', 'dor':'tsedek' }
percs1 = {'likud': 885163, 'yeshatid': 543458, 'avoda': 432118, 'bayityehudi': 345985, 'shas': 331868,
         'yahadut': 195892, 'tenuah': 189167, 'meretz': 172403, 'raamtaal': 138450, 'hadash': 113439,
         'balad': 97030, 'kadima': 78974}
percs1['other'] = 3792742-sum(percs1.values())
total = sum(percs1.values())
#for party in percs:
#    percs[party] /= total

m = Main(0.02)
result1 = sorted(m.perform_bader_ofer(percs1, surplus1).items(), key=lambda x:x[1], reverse = True)
result2 = m.bader_ofer(percs1, surplus1)
for party, mands in result1:
    print (party, mands, result2[party] if party in result2 else 0, mands - result2[party] if party in result2 else mands)

percs2= {'Habayit Hayehudi': 0.0902930187208, 'Yishai': 0.0352718824652, 'Koolanu': 0.0636297779456,
         'Meretz': 0.0405248349533, 'Mahane Tzioni': 0.185939738678, 'Likud': 0.173951818714,
         'Shas': 0.0556348566752, 'Joint': 0.0994893720348,
         'Yahadut Hatorah': 0.0530899600436, 'Yisrael Beitenu': 0.003115436853772964, 'Yesh Atid': 0.0978238584891}
sigma2= {'Habayit Hayehudi': 0.0026946387023503897, 'Yishai': 0.001980196402268915, 'Koolanu': 0.0032029934725094074,
         'Meretz': 0.0015336773868115173, 'Mahane Tzioni': 0.004136856480205219, 'Likud': 0.0033931191177079393,
         'Shas': 0.00249721964962137, 'Joint': 0.0019367576476034803,
         'Yahadut Hatorah': 0.0011624738029221557, 'Yisrael Beitenu': 0.003115436853772964, 'Yesh Atid': 0.00390300273922642}
percs2['other'] = 1 - sum(percs2.values())
surplus2= {'Likud': 'Habayit Hayehudi', 'Shas': 'Yahadut Hatorah', 'Koolanu': 'Yisrael Beitenu', 'Meretz': 'Mahane Tzioni'}
m = Main(0.0325)
percs2b = dict([(party, percs2[party]*120) for party in percs2])
for party, mands in sorted(m.perform_bader_ofer(percs2b, surplus2).items(), key=lambda x:x[1], reverse = True):
    print (party, mands)

surplus2b= {'Likud': 'Habayit Hayehudi', 'Shas': 'Yahadut Hatorah', 'Koolanu': 'Yisrael Beitenu', 'Meretz': 'Joint', 'Yesh Atid':'Mahane Tzioni'}
for party, gains in sorted(m.compare_simulations(percs2, surplus2, surplus2b, sigma2).items()):
    print (party, sum(gains)/len(gains))
    for key, group in itertools.groupby(gains):
        print("-", key, len(list(group))/len(gains))

percs3={ 'Kadima': 758032, 'Likud': 729054, 'YisraelBeitenu':394577, 'Avoda':334900, 'Shas':286300, 'Yahadut':147954, 'RaamTaal':113954, 'IhudLeumi':112570, 'Bayityehudi':96765, 'Hadash':112130, 'Meretz':99611, 'Balad':83739 }
surplus3={'Avoda':'Meretz', 'Likud':'YisraelBeitenu', 'Kadima':'Yerukim', 'Bayityehudi':'IhudLeumi', 'Gil':'Tsabar','Shas':'Yahadut','Hadash':'RaamTaal', 'YisraelHazaka':'YerukaMemad'}
percs3['other'] = 3373490 - sum(percs3.values())
m = Main(0.02)
for party, mands in sorted(m.perform_bader_ofer(percs3, surplus3).items(), key=lambda x:x[1], reverse = True):
    print (party, mands)