from otree.api import *
import random

doc = """
Basic creation of binomial rv.
"""


class C(BaseConstants):
    NAME_IN_URL = 'Rounds'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 2
    r_min = 0
    N_size = 10  # this is the sample size
    P_yellow = [0.75, 0.25]  # yellow and green box respectively
    P_Yellow_Box = 0.90
    P_Yellow_Other = 0.20


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    rand_box = models.FloatField(initial=0.75)
    round_yellow_n = models.IntegerField(initial=5)
    round_green_n = models.IntegerField(initial=5)
    report = models.IntegerField(min=0,
                                 max=100,
                                 label='Please estimate the probability '
                                       'the Yellow box was selected:')
    report_other = models.IntegerField(min=0,
                                       max=100,
                                       label='Please estimate the probability '
                                       'participant 2 selected the Yellow box:')
    round_payoff = models.IntegerField(initial=0)


# PAGES

class New_Round(Page):
    form_model = 'player'

    # after_all_player_arrive = 'set_new_rand_num'
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        import random
        # participant = player.participant
        player.rand_box = random.choices(C.P_yellow, k=1, weights=[C.P_Yellow_Box, 1 - C.P_Yellow_Box])[0]
        marble_sample = random.choices([0, 1], k=C.N_size, weights=[1 - player.rand_box, player.rand_box])
        yellow_n = sum(marble_sample)
        player.round_yellow_n = yellow_n
        player.round_green_n = C.N_size - yellow_n


class RoundSeq(Page):
    form_model = 'player'
    form_fields = ['report', 'report_other']

    @staticmethod
    def js_vars(player):
        # this function generates the sequence of marbles to be pass to JS
        rand_seq = [1] * player.round_yellow_n + [0] * player.round_green_n
        random.shuffle(rand_seq)
        return dict(
            marble_sample=rand_seq,
            N=C.N_size
        )


page_sequence = [
    New_Round,
    RoundSeq
]