from otree.api import *


doc = """
Simple Investment Game
"""


class C(BaseConstants):
    NAME_IN_URL = 'investment_game_test'
    PLAYERS_PER_GROUP = 2
    NUM_ROUNDS = 5
    ENDOWMENT = cu(20)
    MULTIPLIER = 1
    INSTRUCTIONS_TEMPLATE = 'investment_game_test/instructions.html'


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    sent_amount = models.CurrencyField(
        min=cu(0),
        max=cu(1),
        doc="""Investment by P1""",
        label="Do you want to invest 1 point in this round? Please input 0 or 1.",
    )
    sent_back_amount = models.CurrencyField(
        doc="""Returns on investment""",
        label="How much does participant A earns on his or her investment? Please input 0 or 3."
    )


class Player(BasePlayer):
    pass


# FUNCTIONS
def sent_back_amount_choices(group: Group):
    if group.sent_amount == 0:
        return currency_range(0, 0.5, 0.5)
    elif group.sent_amount == 1:
        return currency_range(0, 3, 3)


def set_payoffs(group: Group):
    p1 = group.get_player_by_id(1)
    p2 = group.get_player_by_id(2)
    p1.payoff = C.ENDOWMENT + p1.payoff - group.sent_amount + group.sent_back_amount
    p2.payoff = group.sent_amount * C.MULTIPLIER - group.sent_back_amount


# PAGES
class Introduction(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1

class Send(Page):
    form_model = 'group'
    form_fields = ['sent_amount']

    @staticmethod
    def is_displayed(player: Player):
        return player.id_in_group == 1


class WaitForP1(WaitPage):
    pass


class SendBack(Page):
    form_model = 'group'
    form_fields = ['sent_back_amount']

    @staticmethod
    def is_displayed(player: Player):
        return player.id_in_group == 2

    @staticmethod
    def vars_for_template(player: Player):
        group = player.group

        return dict(tripled_amount=group.sent_amount * C.MULTIPLIER)


class ResultsWaitPage(WaitPage):
    after_all_players_arrive = set_payoffs


class Results(Page):
    pass


page_sequence = [Introduction, Send, WaitForP1, SendBack, ResultsWaitPage, Results]