from otree.api import *



doc = """
Simple trust game
"""


class Constants(BaseConstants):
    name_in_url = 'trust_twp'
    players_per_group = 2
    num_rounds = 1
    endowment = cu(10)
    payoff_cut = cu(4)
    multiplier = 3
    instructions_template = 'trust_with_punishment/instructions.html'


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    sent_amount = models.CurrencyField(
        min=cu(0),
        max=Constants.endowment,
        doc="""Amount sent by P1""",
        label="How much do you want to send to participant B?",
    )
    desired_payback = models.CurrencyField(
        min=cu(0),
        doc="""Amount P1 desires P2 to send back""", 
        label="How much do you desire participant B to send back?"
    )
    payoff_cut = models.CurrencyField(
        choices=[0, Constants.payoff_cut],
        label="How much do you want to cut from B's payoff if they give you less than what you desire?"
    )
    sent_back_amount = models.CurrencyField(
        doc="""Amount sent back by P2""", label="How much do you want to send back?"
    )


class Player(BasePlayer):
    pass


# FUNCTIONS
def sent_back_amount_choices(group: Group):
    return currency_range(cu(0), group.sent_amount * Constants.multiplier, cu(1))


def set_payoffs(group: Group):
    p1 = group.get_player_by_id(1)
    p2 = group.get_player_by_id(2)
    p1.payoff = Constants.endowment - group.sent_amount + group.sent_back_amount
    p2.payoff = Constants.endowment + group.sent_amount * Constants.multiplier - group.sent_back_amount
    if group.sent_back_amount < group.desired_payback:
        p2.payoff -= Constants.payoff_cut


# PAGES
class Send(Page):
    form_model = 'group'
    form_fields = ['sent_amount', 'desired_payback', 'payoff_cut']

    @staticmethod
    def is_displayed(player: Player):
        return player.id_in_group == 1
    
    @staticmethod
    def error_message(player: Player, values):
        if values['desired_payback'] > 3 * values['sent_amount']:
            return "You can not desire more than tripled your sent amount."


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 * Constants.multiplier)


class ResultsWaitPage(WaitPage):
    after_all_players_arrive = set_payoffs


class Results(Page):
    pass


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