from otree.api import *

doc = """
This is a prisoner's dilemma game with side payments as well as symmetric and asymmetric treatments
"""


class Constants(BaseConstants):
    name_in_url = 'sym_game_test'
    players_per_group = 2
    num_rounds = 1

    treatment = ['symmetric', 'asymmetric']

    rich_role = 'Rich'
    poor_role = 'Poor'


class Subsession(BaseSubsession):
    pass


def creating_session(subsession: Subsession):
    import itertools
    treatments = itertools.cycle([True, False])
    for g in subsession.get_groups():
        g.treatment = next(treatments)


class Group(BaseGroup):
    treatment = models.BooleanField(
        doc="Group",
        label="Assigned group",
        choices=[
            [True, 'Asymmetric'],
            [False, 'Symmetric']
        ]
    )
    sp_amount = models.FloatField(
        choices=[0, 1, 2],
        doc="Amount sent by Country A",
        label="How many points do you want to transfer to player B?",
    )
    sp_decision = models.BooleanField(
        blank=True,
        doc="Decision of Country B",
        label="Do you want to accept the transfer?",
        choices=[
            [True, 'Accept'],
            [False, 'Reject'],
        ]
    )


class Player(BasePlayer):
    cooperate = models.BooleanField(
        label="I want to",
        choices=[
            [True, 'contribute the points to the joint project'],
            [False, 'keep the points'],
        ]
    )


def other_player(player: Player):
    return player.get_others_in_group()[0]


def set_payoffs(group: Group):
    p1 = group.get_player_by_role(Constants.poor_role)
    p2 = group.get_player_by_role(Constants.rich_role)
    if group.treatment:
        if group.sp_amount >= 1 and group.field_maybe_none('sp_decision') == 1:
            if p1.cooperate:
                if p2.cooperate:
                    p1.payoff = 4 + group.sp_amount
                    p2.payoff = 10 - group.sp_amount
                else:
                    p1.payoff = 2 + group.sp_amount
                    p2.payoff = 5 - group.sp_amount
            else:
                if p2.cooperate:
                    p1.payoff = 5 + group.sp_amount
                    p2.payoff = 6 - group.sp_amount
                else:
                    p1.payoff = 3 + group.sp_amount
                    p2.payoff = 3 - group.sp_amount
        else:
            if p1.cooperate:
                if p2.cooperate:
                    p1.payoff = 4
                    p2.payoff = 10
                else:
                    p1.payoff = 2
                    p2.payoff = 5
            else:
                if p2.cooperate:
                    p1.payoff = 5
                    p2.payoff = 6
                else:
                    p1.payoff = 3
                    p2.payoff = 3
    else:
        if group.sp_amount >= 1 and group.field_maybe_none('sp_decision') == 1:
            if p1.cooperate:
                if p2.cooperate:
                    p1.payoff = 5 + group.sp_amount
                    p2.payoff = 5 - group.sp_amount
                else:
                    p1.payoff = 2 + group.sp_amount
                    p2.payoff = 6 - group.sp_amount
            else:
                if p2.cooperate:
                    p1.payoff = 6 + group.sp_amount
                    p2.payoff = 2 - group.sp_amount
                else:
                    p1.payoff = 3 + group.sp_amount
                    p2.payoff = 3 - group.sp_amount
        else:
            if p1.cooperate:
                if p2.cooperate:
                    p1.payoff = 5
                    p2.payoff = 5
                else:
                    p1.payoff = 2
                    p2.payoff = 6
            else:
                if p2.cooperate:
                    p1.payoff = 6
                    p2.payoff = 2
                else:
                    p1.payoff = 3
                    p2.payoff = 3


# PAGES


class Roles(Page):
    pass


class Sidepayments(Page):
    form_model = 'group'
    form_fields = ['sp_amount']

    @staticmethod
    def is_displayed(player: Player):
        return player.role == Constants.rich_role


class WaitforOffer(WaitPage):

    @staticmethod
    def is_displayed(player: Player):
        return player.role == Constants.poor_role


class Decision(Page):

    @staticmethod
    def is_displayed(player: Player):
        return player.role == Constants.poor_role

    form_model = 'group'
    form_fields = ['sp_decision']


class Contribution(Page):
    form_model = 'player'
    form_fields = ['cooperate']


class ResultsWaitPage(WaitPage):
    after_all_players_arrive = set_payoffs


class Results(Page):
    @staticmethod
    def vars_for_template(player: Player):
        opponent = other_player(player)
        return dict(
            opponent=opponent,
            same_choice=player.cooperate == opponent.cooperate,
            my_decision=player.field_display('cooperate'),
            opponent_decision=opponent.field_display('cooperate'),
            opponent_payoff=opponent.payoff,
        )


page_sequence = [Roles, Sidepayments, WaitforOffer, Decision, Contribution,
                 ResultsWaitPage,
                 Results]