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 = 3
    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)
    report_other = models.IntegerField(min=0,
                                       max=100)
    round_payoff = models.IntegerField(initial=0)
    payoff_final = models.CurrencyField(initial=0)


# PAGES

class New_Round(Page):
    form_model = 'player'

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        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
        )

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        prob_win = (player.rand_box == 0.75) * (1 - (1 - player.report/100) ** 2) \
                   + (player.rand_box == 0.25) * (1 - (player.report/100) ** 2)
        bet_result = random.choices([4, 0], weights=[prob_win, 1 - prob_win], k=1)[0]
        expected_diff = player.report_other/100 - 0.5  # I have to include here the prediction. Probably from a csv file
        bet_other_result = random.choices([4, 0], k=1,
                                          weights=[1 - expected_diff ** 2, expected_diff ** 2])[0]
        player.payoff = bet_result + bet_other_result

        participant = player.participant
        # if it's the last round
        if player.round_number == C.NUM_ROUNDS:
            random_round = random.randint(1, C.NUM_ROUNDS)
            participant.selected_round = random_round
            player_in_selected_round = player.in_round(random_round)
            player.payoff_final = cu(player_in_selected_round.payoff)
            participant.payoff = player.payoff_final


class Results(Page):
    form_model = 'player'


page_sequence = [
    New_Round,
    RoundSeq,
    # Results
]