from otree.api import *



class C(BaseConstants):
    NAME_IN_URL = 'public_goods_large'
    PLAYERS_PER_GROUP = 10
    NUM_ROUNDS = 6
    ENDOWMENT = cu(20)
    MULTIPLIER = 1.6
    MAX_PUNISHMENT = 10
    INSTRUCTIONS_TEMPLATE = 'public_goods_large/instructions.html'
    PUNISHMENT_SCHEDULE = {
        0: 0,
        1: 1,
        2: 2,
        3: 4,
        4: 6,
        5: 9,
        6: 12,
        7: 16,
        8: 20,
        9: 25,
        10: 30,
    }


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    total_contribution = models.CurrencyField()
    individual_share = models.CurrencyField()


def make_punishment_field(id_in_group):
    return models.IntegerField(
        min=0, max=C.MAX_PUNISHMENT, label="Punishment to player {}".format(id_in_group)
    )

class Player(BasePlayer):
    contribution = models.CurrencyField(
        min=0, max=C.ENDOWMENT, label="How much will you contribute?"
    )
    punish_p1 = make_punishment_field(1)
    punish_p2 = make_punishment_field(2)
    punish_p3 = make_punishment_field(3)
    punish_p4 = make_punishment_field(4)
    punish_p5 = make_punishment_field(5)
    punish_p6 = make_punishment_field(6)
    punish_p7 = make_punishment_field(7)
    punish_p8 = make_punishment_field(8)
    punish_p9 = make_punishment_field(9)
    punish_p10 = make_punishment_field(10)
    cost_of_punishing = models.CurrencyField()
    punishment_received = models.CurrencyField()


# FUNCTIONS
def creating_session(subsession):
    subsession.group_randomly()

def get_self_field(player: Player):
    return 'punish_p{}'.format(player.id_in_group)

def punishment_fields(player: Player):
    return ['punish_p{}'.format(p.id_in_group) for p in player.get_others_in_group()]

def set_payoffs(group: Group):
    players = group.get_players()
    contributions = [p.contribution for p in players]
    group.total_contribution = sum(contributions)
    group.individual_share = (
        group.total_contribution * C.MULTIPLIER / C.PLAYERS_PER_GROUP
    )
    for p in players:
        payoff_before_punishment = C.ENDOWMENT - p.contribution + group.individual_share
        self_field = get_self_field(p)
        punishments_received = [getattr(other, self_field) for other in p.get_others_in_group()]
        p.punishment_received = min(10, sum(punishments_received))
        punishments_sent = [getattr(p, field) for field in punishment_fields(p)]
        p.cost_of_punishing = sum(C.PUNISHMENT_SCHEDULE[points] for points in punishments_sent)
        p.payoff = payoff_before_punishment * (1 - p.punishment_received / 10) - p.cost_of_punishing

def vars_for_admin_report(subsession):
    all_contributions = []
    all_payoffs = []
    for subsession in subsession.in_all_rounds():
        for group in subsession.get_groups():
            for p in group.get_players():
                contributions = p.contribution
                all_contributions.append(contributions)
                payoffs = p.payoff
                all_payoffs.append(payoffs)

    return dict(all_payoffs=all_payoffs, all_contributions=all_contributions)

# PAGES
class Introduction(Page):
    pass

class Contribute(Page):
    form_model = 'player'
    form_fields = ['contribution']

class WaitPage1(WaitPage):
    pass

class Punish(Page):
    form_model = 'player'
    get_form_fields = punishment_fields

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            other_players=player.get_others_in_group(), schedule=C.PUNISHMENT_SCHEDULE.items(),
        )

class WaitPage2(WaitPage):
    after_all_players_arrive = set_payoffs


class Results(Page):
    pass


page_sequence = [Introduction, Contribute, WaitPage1, Punish, WaitPage2, Results]