from otree.api import (
    models, widgets, BaseConstants, BaseSubsession, BaseGroup, BasePlayer,
    Currency as c, currency_range
)

doc = ''
class Constants(BaseConstants):
    name_in_url = 'game1'
    players_per_group = None
    num_rounds = 100
    num_sites = 4
    rhino_distribution = (0.8, 0.3, 0.8, 0.3)
class Subsession(BaseSubsession):
    pass
class Group(BaseGroup):
    pass
class Player(BasePlayer):
    poacher_site = models.IntegerField(choices=[[1, 'Site 1'], [2, 'Site 2'], [3, 'Site 3'], [4, 'Site 4']], label='Which site would you like to visit for this round?', widget=widgets.RadioSelect)
    ranger_site = models.IntegerField()
    def get_payoff(self):
        import random
        self.get_ranger_site()
        i = self.poacher_site
        j = self.participant.vars["ranger_site"]
        r = random.random() < Constants.rhino_distribution[i-1]
        if i!=j and r:
            self.payoff = 1
        elif i == j:
            self.payoff = -1
        else:
            self.payoff = 0
        self.update_ranks()
    def get_ranger_site(self):
        import random
        if random.random() < 0.95:
            self.participant.vars["ranger_site"] = self.utility_match()
        else:
            self.participant.vars["ranger_site"] = self.probability_match() # randomness
    def update_ranks(self):
        import random
        j = self.participant.vars["ranger_site"]-1
        ranks = self.participant.vars["ranks"]
        d = (ranks[j]-0.5) / 100
        u = 1-d
        if self.payoff == -1:
            if random.random() < u and ranks[j] < 100:
                ranks[j] += 1
        else:
            if random.random() < d and ranks[j] > 0:
                ranks[j] -= 1  
        self.participant.vars["ranks"] = ranks
    def get_bonus(self):
        return max(round(int(self.participant.payoff)*0.1, 2),0)
    def probability_match(self):
        import random
        r = [0] * Constants.num_sites
        s = sum(Constants.rhino_distribution)
        for i in range(Constants.num_sites):
            r[i] = Constants.rhino_distribution[i]/s
        for i in range(1, Constants.num_sites):
            r[i] += r[i-1]
        u = random.random()
        for i in range(Constants.num_sites):
            if u <= r[i]:
                return i 
    def utility_match(self):
        import random
        # calculate utility
        utility = [0] * Constants.num_sites
        r = [0] * Constants.num_sites
        for i in range(Constants.num_sites):
            r[i] = self.participant.vars["ranks"][i]/100
        for i in range(Constants.num_sites):
            utility[i] = r[i] * (1+Constants.rhino_distribution[i])
        max_utility = max(utility)
        idx = []
        for i in range(Constants.num_sites):
            if max_utility > 0 and utility[i] >= 0.75*max_utility:
                idx.append(i)
        if not idx:
            return self.probability_match()+1
        else:
            est = [utility[site] for site in idx]
            # normalize
            s = sum(est)
            for i in range(len(est)):
                est[i] = est[i] / s
            for i in range(1, len(est)):
                est[i] += est[i-1]
            # utility matching
            u = random.random()
            for i in range(len(est)):
                if u <= est[i]:
                    return idx[i]+1