from otree.api import *
import random

doc = """
Your app description
"""


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

    show_up = 125
    b_bonus = 250


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    treatment = models.StringField()  # T1, T1_b, T2, T2_b
    type = models.StringField()  # A, B

    decision_A = models.StringField(choices=['A1', 'A2'])
    decision_B = models.IntegerField()


def creating_session(subsession):
    if subsession.session.vars.get('group_matrix'):
        group_matrix = subsession.session.vars['group_matrix']
        subsession.set_group_matrix(group_matrix)

    for player in subsession.get_players():
        # set treatment from session config, if not set before
        if player.participant.vars.get('treatment') is None:
            player.treatment = subsession.session.config['treatment']
            player.participant.vars['treatment'] = player.treatment
        else:
            player.treatment = player.participant.vars['treatment']

        # set types by group ids, if not grouped before
        if player.participant.vars.get('type') is None:
            if player.id_in_group == 1:
                player.type = "A"
            if player.id_in_group == 2:
                player.type = "B"
            player.participant.vars['type'] = player.type
        else:
            player.type = player.participant.vars['type']


# PAGES
class Instructions(Page):
    pass


class Roles(Page):
    pass


class DecisionA(Page):
    form_model = 'player'
    form_fields = ['decision_A']

    def is_displayed(player: Player):
        return player.type == "A" and (player.treatment == "T1" or player.treatment == "T2")


class DecisionB(Page):
    form_model = 'player'
    form_fields = ['decision_B']

    def is_displayed(player: Player):
        return player.type == "B"


class WaitForAllPage(WaitPage):
    wait_for_all_groups = True


class PlayerAWaitPage(WaitPage):
    body_text = "Please wait while Person Bs make their choices in Part 1"

    def is_displayed(player: Player):
        return player.type == "A" and (player.treatment == "T1_B" or player.treatment == "T2_B")


class WaitEndOfPart1(WaitPage):
    def after_all_players_arrive(group: Group):
        players = group.get_players()
        player_A = [p for p in players if p.type == "A"][0]
        player_B = [p for p in players if p.type == "B"][0]
        # randomize answer of A if its a base treatment
        treatment = players[0].treatment
        if treatment == 'T1_B' or treatment == 'T2_B':
            player_A.decision_A = random.choice(['A1', 'A2'])

        # save decisions for recall in p2 to participants
        decision_A = player_A.decision_A
        decision_B = player_B.decision_B
        for p in players:
            p.participant.vars['p1'] = {
                'decision_A': decision_A,
                'decision_B': decision_B,
            }

        # calculate payoff for players
        if player_A.decision_A == "A2":
            payoff_A = Constants.show_up + player_B.decision_B
            payoff_B = Constants.b_bonus + Constants.show_up - player_B.decision_B
        else:
            payoff_A = Constants.show_up
            payoff_B = Constants.b_bonus + Constants.show_up
        player_A.participant.vars['payoffs'] = {
            'p1': payoff_A,
            'p1_b': payoff_B,
            'p1_a': payoff_A
        }
        player_B.participant.vars['payoffs'] = {
            'p1': payoff_B,
            'p1_a': payoff_A,
            'p1_b': payoff_B
        }

        print("decisions saved to participants! payoffs calculated")


class Completion(Page):
    def vars_for_template(player: Player):
        return {
            'pvars': player.participant.vars
        }


page_sequence = [Instructions,
                 WaitForAllPage,
                 Roles,
                 WaitForAllPage,
                 DecisionA,
                 DecisionB,
                 PlayerAWaitPage,
                 WaitEndOfPart1,
                 Completion]