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

doc = '7-rounds baseline investment game'
class Constants(BaseConstants):
    name_in_url = 'Prova_base'
    players_per_group = 3
    num_rounds = 4
    alpha = 5
    beta = 10
    gamma = 3
    teta = 1
    endowment = c(7)
    amounttax = 3
def creating_session(subsession):
    session = subsession.session
    import itertools
    treatments = itertools.cycle(['baseline','baseline','baseline','information', 'information', 'information','tax','tax','tax'])
    for p in subsession.get_players():
        p.treatment = next(treatments)
    
    
    if subsession.round_number == 1:
        players = subsession.get_players()
        baseline_players = [p for p in players if p.treatment == 'baseline']
        information_players = [p for p in players if p.treatment == 'information']
        tax_players = [p for p in players if p.treatment == 'tax']
        group_matrix = []
    
                # pop elements from M_players until it's empty
        while baseline_players:
            new_group1 = [
                        baseline_players.pop(),
                        baseline_players.pop(),
                        baseline_players.pop(),
                    ]
            new_group2 = [
                        information_players.pop(),
                        information_players.pop(),
                        information_players.pop(),
            ]
            new_group3 = [
                        tax_players.pop(),
                        tax_players.pop(),
                        tax_players.pop(),
            ]
    
            group_matrix.append(new_group1)
            group_matrix.append(new_group2)
            group_matrix.append(new_group3)
        print('matrice gruppi',group_matrix)
        subsession.set_group_matrix(group_matrix)
    
    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()
    calculation = models.CurrencyField()
    treatment = models.StringField()
    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()
    treatment = models.StringField()