from otree.api import *

class C(BaseConstants):
    NAME_IN_URL = 'prisoners_dilemma'
    PLAYERS_PER_GROUP = 2
    NUM_ROUNDS = 2

    PUSH = 'Push'
    PULL = 'Pull'

class Subsession(BaseSubsession):
    def creating_session(self):
        self.group_randomly()

class Group(BaseGroup):

    def set_payoffs(self):
        p1 = self.get_players()[0]
        p2 = self.get_players()[1]

        # Both push
        if p1.decision == C.PUSH and p2.decision == C.PUSH:
            p1.payoff = 400
            p2.payoff = 400

        # P1 push, P2 pull
        elif p1.decision == C.PUSH and p2.decision == C.PULL:
            p1.payoff = 0
            p2.payoff = 700

        # P1 pull, P2 push
        elif p1.decision == C.PULL and p2.decision == C.PUSH:
            p1.payoff = 700
            p2.payoff = 0

        # Both pull
        elif p1.decision == C.PULL and p2.decision == C.PULL:
            p1.payoff = 300
            p2.payoff = 300

class Player(BasePlayer):
    prisoners_passcode = models.StringField(blank=True)
    decision = models.StringField(
        choices=[C.PUSH, C.PULL],
        widget=widgets.RadioSelect,
        label="Choose your action:"
    )

class Introduction(Page):
    form_model = 'player'
    form_fields = ['prisoners_passcode']

    def is_displayed(player):
        return player.round_number == 1

class Decision(Page):
    form_model = 'player'
    form_fields = ['decision']


class ResultsWaitPage(WaitPage):
    def after_all_players_arrive(self):
        self.group.set_payoffs()


class Results(Page):

    def vars_for_template(player):
        opponent = player.get_others_in_group()[0]

        return dict(
            my_decision=player.decision,
            opponent_decision=opponent.decision,
            my_payoff=player.payoff,
            opponent_payoff=opponent.payoff,
        )


class BetweenRounds(Page):
    def vars_for_template(player):
        return {
            'last_round_payoff': player.payoff,
            'round_number': player.round_number,
        }

    def is_displayed(player):
        return player.round_number == 1

page_sequence = [
    Introduction,
    Decision,
    ResultsWaitPage,
    Results,
    BetweenRounds
]