import random, itertools
import numpy as np
from otree.api import *
c = cu

doc = """"""
# CLASSES
class C(BaseConstants):
    NAME_IN_URL = 'example'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 1
    BANDIT_ARMS = 5
    BANDIT_MU = [10, 15] # the 2 grand means, use this to generate means of options
    BANDIT_SD = 0.5
    SCALE = 2

class Subsession(BaseSubsession):
    pass

class Group(BaseGroup):
    pass

class Player(BasePlayer):
    value = models.FloatField() # record the current sampled value
    arm = models.IntegerField() # record the current sampled arm
    seq_value = models.LongStringField(initial= "")
    seq_arm = models.LongStringField(initial= "")
    choice = models.IntegerField(label= "Choose a restaurant to invest in:", 
                                 choices= range(1,C.BANDIT_ARMS+1))
    estimation = models.IntegerField(
        label= "How many customers do you expect to vist your restaurant on average?")
    chosen_mu = models.FloatField()
    drawn = models.FloatField() # record the drawn value from the chosen option
    is_understand = models.BooleanField(initial=True)
    # FUNCTIONS
    def make_field(label,answer):
        return models.IntegerField(
            blank= False,
            label = label,
            widget = widgets.RadioSelect,
            choices= [
                [1, answer[0]],
                [2, answer[1]],
                [3, answer[2]],
                [4, answer[3]]]
        )
    q1 = make_field(label="There are 100 restaurants in city A. Based on the national report, how many restaurants in this city are busier than the rest?",
                    answer= ['A. 15','B. 10','C. 20','D. 30'])
    q2 = make_field(label="My reward in this task is determined by:",
                    answer= ['A. The number of customers visited my invested restaurant in 1 random day',
                             'B. The number of customers visited the best restaurant in the city',
                             'C. The average number of customers visited my invested restaurant',
                             'D. The smallest number of customers visited my invested restaurant'])
    q3 = make_field(label="The observed customers count of a restaurant in 5 random days are: 30, 40, 40, 30, 20. Which number below is closest to the average of this number sequence?",
                    answer= ['A. 10','B. 20','C. 30','D. 60'])

   
# Functions

def custom_export(players):
    yield ['participant_code', 'id_in_group'] # THESE ARE 2 OBJECTS OF PLAYERS
    for p in players:
        pp = p.participant
        yield [pp.code, p.id_in_group]

# PAGES
class Introduction(Page):
    pass

class Bandit(Page):
    form_model = 'player'
    form_fields = ['choice','estimation']
    @staticmethod
    def live_method(player: Player, data): # receive data from the page         
        # generate the options values using player id 
        temp = list(np.random.normal(C.BANDIT_MU[1],C.BANDIT_SD,int(1/5*C.BANDIT_ARMS))) + list(np.random.normal(C.BANDIT_MU[0],C.BANDIT_SD, int(4/5*C.BANDIT_ARMS)))
        np.random.shuffle(temp)
        Bandit.options_mu = [ round(i,2) for i in temp ] 
        options_mu = Bandit.options_mu
        print(Bandit.options_mu)
        
        # receive the button press and sample the according distribution
        player.arm = data
        mu = options_mu[int(data) -1] * C.SCALE #C.BANDIT_MU[int(data)-1]
        diff = 5*C.SCALE
        player.value = round(random.uniform(mu - diff, mu + diff))
        player.seq_value += ";" + str(player.value)
        player.seq_arm += ";" + str(player.arm)
        response = dict(arm = player.arm, value = player.value)
        
        # create objects for feedback
        player.drawn = round(random.uniform(mu - diff, mu + diff)) # show in feedback
        player.chosen_mu = mu # to call from feedback
        
        return {player.id_in_group: response}
    
    @staticmethod
    def js_vars(player): # to include set size variable in javascript
        return dict(setSize = C.BANDIT_ARMS)    
    
    @staticmethod
    def error_message(player, values):
        if player.seq_arm == "":
            return "You haven't sampled any restaurant"


class Quiz(Page):
    form_model = 'player'
    form_fields = ['q1','q2','q3']
    @staticmethod
    def error_message(player: Player, values):
        solutions = dict(q1=3, q2=3, q3=3)

        if values != solutions:
            player.is_understand = False

class Exclude(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.is_understand == False
        
page_sequence = [Introduction, Bandit, Quiz, Exclude] # this is the easiet way to have 4 repeated trials, we need to develop an feedback app and an Result app