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

doc = '5-rounds baseline2 investment game'
class Constants(BaseConstants):
    name_in_url = 'Baseline2_InvGame'
    players_per_group = 3
    num_rounds = 5
    alpha = 5
    beta = 10
    gamma = 3
    teta = 1
    endowment = c(7)
def creating_session(subsession):
    session = subsession.session
    if subsession.round_number == 1:
        subsession.group_randomly()
    else:
        subsession.group_like_round(1)
class Subsession(BaseSubsession):
    creating_session = creating_session
def set_payoff(group):
    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()
    set_payoff = set_payoff
class Player(BasePlayer):
    investment = models.CurrencyField(label='How many points do you allocate in the RED investment? The remaining will be allocated in the WHITE investment.', max=Constants.endowment, min=0)
    white_investment = models.CurrencyField()