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

doc = '7-rounds treatment investment game'
class Constants(BaseConstants):
    name_in_url = 'Treatment_InvGameRESD1'
    players_per_group = 3
    num_rounds = 5
    alpha = 5
    beta = 10
    gamma = 3
    teta = 1
    endowment = c(7)
    amounttax = 3
def creating_session(subsession):
    session = subsession.session
    #Balanced treatment groups
    import itertools
    treatments = itertools.cycle(['baseline','tax','information'])
    for g in subsession.get_groups():
        g.treatment = next(treatments)
class Subsession(BaseSubsession):
    creating_session = creating_session
def set_payoff_tax(group):
    if group.treatment == 'tax':
        players = group.get_players()
        investments = [p.investment for p in players]
        group.total_investment = sum(investments)
        for p in players:
            calculation = Constants.alpha * (Constants.endowment - p.investment) - (21 - group.total_investment) + (Constants.beta - Constants.amounttax) * p.investment - Constants.gamma * group.total_investment + group.total_investment * Constants.amounttax / Constants.players_per_group
            if (calculation > 0):
                p.payoff = calculation
            else:
                p.payoff == 0
            p.white_investment = ( Constants.endowment - p.investment)
        payoffs = [p.payoff for p in players]
        group.total_payoff = sum(payoffs)
        group.average_payoff = group.total_payoff / Constants.players_per_group
    else:
        players = group.get_players()
        investments = [p.investment for p in players]
        group.total_investment = sum(investments)
        for p in players:
            calculation = Constants.alpha * (Constants.endowment - p.investment) - (21 - group.total_investment) + Constants.beta * p.investment - Constants.gamma * group.total_investment 
            if (calculation > 0):
                p.payoff = calculation
            else:
                p.payoff == 0
            p.white_investment = ( Constants.endowment - p.investment)
        payoffs = [p.payoff for p in players]
        group.total_payoff = sum(payoffs)
        group.average_payoff = group.total_payoff / Constants.players_per_group
class Group(BaseGroup):
    total_investment = models.CurrencyField()
    total_payoff = models.CurrencyField()
    average_payoff = models.CurrencyField()
    treatment = models.StringField()
    set_payoff_tax = set_payoff_tax
class Player(BasePlayer):
    investment = models.CurrencyField(label='How many points do you allocate in the RED investment? The rest will be allocated in the WHITE investment.', max=Constants.endowment, min=0)
    white_investment = models.CurrencyField()