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

doc = """"""
# CLASSES
class C(BaseConstants):
    NAME_IN_URL = 'block2'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 2
    SETSIZE = 30
    # grand distributions
    MU = [2,8]
    IS_CORPEV_POS = [1,0] # 1 ev increase with p, 0 ev neg corr with p
    # feedback about performance
    TIME = 7
    FIXED_EURO = np.round((TIME*0.1)/ 0.86,1)
    PAYRATE = 0.0009#(1.25*FIXED_EURO)/ (170*8 + EST_CORRECT*8) # minimum pay rate

class Subsession(BaseSubsession):
    pass

class Group(BaseGroup):
    pass

class Player(BasePlayer):
    is_corpev_pos = models.IntegerField()
    seq_value = models.LongStringField(initial= "")
    seq_arm = models.LongStringField(initial= "")
    #setSize = models.IntegerField()
    choice = models.IntegerField(label= "Choose a slot machine to play for real:", 
                                 choices= range(1,C.SETSIZE+1))
    estimation = models.IntegerField(
        label= "How many times do you expect to win if you get 100 draws from this slot machine?",  min = 0, max = 100)
    chosen_x = models.FloatField()
    chosen_p = models.FloatField()
    
    allXP = models.LongStringField()
    
    # demographic questions
    age = models.IntegerField(label='What is your year of birth?', min=1923, max=2004)
    gender = models.StringField(
        choices=[['Male', 'Male'], ['Female', 'Female'], ['Other', 'Other']],
        label='What is your gender?',
        widget=widgets.RadioSelect)
    feedback = models.LongStringField(
        label="Please describe your search strategy in the task, and how you decide which restaurant is the best:", 
        initial= "")
    comment = models.LongStringField(label= "What can we do to improve the experience of future participants?")

    
# Functions
def choice_choices(player):
    choices = range(1, C.SETSIZE+1)
    return choices
def age_choices(player):
    ages = range(1923, 2006)
    return ages

def creating_session(subsession):
    # balanced treatment for within subject is_corpev_pos
    orderTreatment = list(itertools.permutations(C.IS_CORPEV_POS.copy()))
    treatments = itertools.cycle(itertools.product(orderTreatment))
    for player in subsession.get_players():
        t = next(treatments) # t contain a list of shuffled is_corpev_pos
        participant = player.participant
        participant.is_corpev_pos = t[0]
        player.is_corpev_pos = participant.is_corpev_pos[subsession.round_number-1]
    
    for player in subsession.get_players():
        np.random.seed(player.id_in_subsession + 3)
        n = C.SETSIZE
        # is_corpev_pos == 1, ev increase with p
        ev1 = np.round(np.linspace(C.MU[0],C.MU[1], num = n),2)
        p1 = np.round(np.linspace(1/(n+1),n/(n+1), num = n),2)
        x1 = np.round(ev1/p1,2)
  
        # generate p and x from sds, similar sd, different ev
        p2 = np.round(np.linspace(n/(n+1),1/(n+1), num = n),2)
        x2 = np.round(ev1/p2,2)

        for player in subsession.get_players():
            if player.is_corpev_pos == 1:
                pairs = list(zip(x1,p1)) 
                random.shuffle(pairs)
                shuffled_x, shuffled_p = zip(*pairs)
            elif player.is_corpev_pos == 0:
                pairs = list(zip(x2,p2)) 
                random.shuffle(pairs)
                shuffled_x, shuffled_p = zip(*pairs)
            # create dictionary
            data = {'x':shuffled_x , 'p': shuffled_p }
            # convert dictionary to json string
            xPText = json.dumps(data)
            player.allXP = xPText
                

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, data): # receive data from the page
        player.seq_arm = data['seq_arm']
        player.seq_value = data['seq_value']
    # find the true mean of chosen option and calculate payoff    
    @staticmethod
    def before_next_page(player, timeout_happened):
        # calculate whether the estimation within 10% mean
        player.chosen_x = json.loads(player.allXP)['x'][player.choice -1]
        player.chosen_p = json.loads(player.allXP)['p'][player.choice -1]
    
    @staticmethod
    def js_vars(player): # to include set size variable in javascript
        return dict(setSize = C.SETSIZE,
                    allXP = player.allXP)    
    @staticmethod
    def error_message(player, values):
        if player.seq_arm == "":
            return "You haven't sampled any option"
        
class Demographic(Page):
    form_model = 'player'
    form_fields = ['age', 'gender','feedback','comment']
    @staticmethod
    def is_displayed(player: Player):      
        return player.round_number == 2

class Ending(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):      
        return player.round_number == 2
    

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