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

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

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= "")
    setSize = models.IntegerField()
    scale = models.IntegerField()
    choice = models.IntegerField(label= "Choose a restaurant to invest in:")
    estimation = models.IntegerField(
        label= "How many customers do you expect to vist your restaurant on average?")
    is_feedback = models.BooleanField()
    chosen_mu = models.FloatField()
    drawn = models.FloatField() # record the drawn value from the chosen option
    allMus = models.LongStringField()
   
# Functions
def choice_choices(player):
    choices = range(1, player.setSize+1)
    return choices
    
def creating_session(subsession): # this will be call after each round
    if subsession.round_number == 1:
        for player in subsession.get_players():
            participant = player.participant
            participant.is_feedback = random.choice([0,1])
    for player in subsession.get_players():
        player.is_feedback = player.participant.is_feedback 
    if subsession.round_number == 1:
        for player in subsession.get_players():
            participant = player.participant
            participant.setSize = C.BANDIT_ARMS.copy()
            participant.scale = C.SCALE.copy()
            random.shuffle(participant.setSize)
            random.shuffle(participant.scale)
    for player in subsession.get_players():
        player.scale = player.participant.scale[subsession.round_number-1]
        player.setSize = player.participant.setSize[subsession.round_number-1]



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):
    @staticmethod
    def is_displayed(player: Player):
        return player.round_number == 1
    
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 
        np.random.seed(player.id_in_subsession) # HAVE TO CHANGE THIS TO PROLIFIC ID OR AN UNIQUE PARTICIPANT IDENTIFIER
        temp = list(np.random.normal(C.BANDIT_MU[1],C.BANDIT_SD,int(1/5*player.setSize))) + list(np.random.normal(C.BANDIT_MU[0],C.BANDIT_SD, int(4/5*player.setSize)))
        np.random.seed(player.id_in_subsession) # HAVE TO CHANGE THIS TO PROLIFIC ID OR AN UNIQUE PARTICIPANT IDENTIFIER
        np.random.shuffle(temp)
        Bandit.options_mu = [ round(i,2) for i in temp ] 
        options_mu = Bandit.options_mu
        player.participant.options_mu = options_mu
        print(Bandit.options_mu)
        print(player.id_in_subsession)
        print(options_mu)
        # receive the button press and sample the according distribution
        player.arm = data
        mu = options_mu[int(data) -1] * player.scale #C.BANDIT_MU[int(data)-1]
        diff = 5*player.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.allMus = ";".join(str(l) for l in Bandit.options_mu)
        return {player.id_in_group: response}
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        player.chosen_mu = round(Bandit.options_mu[player.choice -1] * player.scale)
        player.payoff = (player.chosen_mu /player.scale)* 0.01
        
    @staticmethod
    def js_vars(player): # to include set size variable in javascript
        return dict(setSize = player.setSize)    
    
    @staticmethod
    def error_message(player, values):
        if player.seq_arm == "":
            return "You have to sample at least 1 machine"

class Feedback(Page):
    @staticmethod
    def is_displayed(player):        
        return (player.is_feedback == 1 and player.round_number == 4)
    @staticmethod   
    def js_vars(player): # to include set size variable in javascript
        round = list(range(1,C.NUM_ROUNDS+1))
        fullSetSize = []
        choice = []
        drawn = []
        estimation = []
        mu = []
        for r in round:
            fullSetSize.append(player.in_round(r).setSize)
            choice.append(player.in_round(r).choice)
            drawn.append(player.in_round(r).drawn)
            estimation.append(player.in_round(r).estimation)  
            mu.append(player.in_round(r).chosen_mu)
        return dict(round = round, 
                    fullSetSize = fullSetSize,
                    choice = choice,
                    drawn = drawn,
                    estimation = estimation,
                    mu = mu)

page_sequence = [Introduction, Bandit,Feedback] # this is the easiet way to have 4 repeated trials, we need to develop an feedback app and an Result app