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


author = 'Your name here'

doc = """
10 rounds, partner matching, pay for all rounds, 2 treatments
"""


class Constants(BaseConstants):
    name_in_url = 'pgg_punish_ext'
    players_per_group = 3
    num_rounds = 10

    endow = c(100)
    multiplier = 1.5

class Subsession(BaseSubsession):
    def creating_session(self):
        if self.round_number == 1:
            import itertools
            treatments = itertools.cycle([True, False])
            for g in self.get_groups():
                g.punishment = next(treatments)
                for p in g.get_players():
                    p.participant.vars['treatment'] = g.punishment
        else:
            for g in self.get_groups():
                g.punishment = g.get_player_by_id(1).participant.vars['treatment']


class Group(BaseGroup):
    punishment = models.BooleanField() #Is True if punishment treatment
    total_contribution = models.CurrencyField()
    individual_share = models.CurrencyField()

    def set_payoffs(self):
        players = self.get_players()
        contributions = [p.contribution for p in players]
        self.total_contribution = sum(contributions)
        self.individual_share = self.total_contribution * Constants.multiplier / Constants.players_per_group
        if self.punishment:
            if self.get_player_by_id(1).punish1: #Returns the player in the group with the given ID and reads the punish1 field of this player
                self.get_player_by_id(2).punished += 1 #adds 1 to the punishment counter of the respective person
                self.get_player_by_id(1).total_punish += 1
            if self.get_player_by_id(1).punish2:
                self.get_player_by_id(3).punished += 1
                self.get_player_by_id(1).total_punish += 1
            if self.get_player_by_id(2).punish1:
                self.get_player_by_id(1).punished += 1
                self.get_player_by_id(2).total_punish += 1
            if self.get_player_by_id(2).punish2:
                self.get_player_by_id(3).punished += 1
                self.get_player_by_id(2).total_punish += 1
            if self.get_player_by_id(3).punish1:
                self.get_player_by_id(1).punished += 1
                self.get_player_by_id(3).total_punish += 1
            if self.get_player_by_id(3).punish2:
                self.get_player_by_id(2).punished += 1
                self.get_player_by_id(3).total_punish += 1
            for p in players:
                p.payoff = Constants.endow - p.contribution + self.individual_share - c(2)*p.total_punish - c(10)*p.punished
        else:
            for p in players:
                p.payoff = Constants.endow - p.contribution + self.individual_share


class Player(BasePlayer):
    contribution = models.CurrencyField(label="How much do you want to contribute?",
                                       min=0, max=Constants.endow)

    punished = models.IntegerField(initial=0)  #by how many?

    total_punish = models.IntegerField(initial=0)

    punish1 = models.BooleanField(label="Do you want to punish this player?",
                                  choices=[[True, "Yes"], [False, "No"]],
                                  widget=widgets.RadioSelect)

    punish2 = models.BooleanField(label="Do you want to punish this player?",
                                  choices=[[True, "Yes"], [False, "No"]],
                                  widget=widgets.RadioSelect)