from otree.api import (
    models, widgets, BaseConstants, BaseSubsession, BaseGroup, BasePlayer,
    Currency as c, currency_range
)

import numpy as np
import random


author = 'Your name here'

doc = """
Multiarmed bandit problems for studying trade-offs between exploration and exploitation
"""


class Constants(BaseConstants):
    name_in_url = 'multiarmed_bandit'
    players_per_group = None
    history_template = 'multiarmed_bandit//History.html'
    history_previous_template = 'multiarmed_bandit//History_previous.html'
    cost = 10

    # info_round = 0 ## new information is added after info_round

    interaction_length = [20,30]
    round_in_interactions = [i + 1 for length in interaction_length for i in range(length)]
    interactions = [i + 1 for i, length in enumerate(interaction_length) for j in range(length)]

    # pparams80 = np.array([[24,  .8], [40,  .8], [40,  .4], [56,  .8], [24,  .4]])
    # pparams80t = np.array([[64,  .8], [40,  .8], [40,  .4], [56,  .8], [24,  .4]])
    # pparams = np.array([[50, .10], [30, .10],  [30,  .5], [50,  .5], [70, .10]])
    # pparamst = np.array([[50, .10], [80, .10], [30,  .5], [50,  .5], [70, .10]])
    # pparams120 = np.array([[60, .12], [60, .6], [84, .12], [36, .12], [36, .6]])
    # pparams120t = np.array([[60, .12], [60, .6], [84, .12], [96, .12], [36, .6]])
    #
    # info_round = 0 ## new information is added after info_round
    # interaction_length = [6,6,6]
    # round_in_interactions = np.tile(np.arange(6)+1, reps=3)
    # interactions = np.tile(np.arange(3) + 1, reps=(6, 1)).T.flatten()

    num_rounds = sum(interaction_length) # change num_rounds for testing purpose, but need to make sure that number_sequence


class Subsession(BaseSubsession):
    def creating_session(self):
        # this is run before the start of every round
        players = self.get_players()
        num_players = len(players)
        ppg = int(num_players/4)
        round_in_interaction = Constants.round_in_interactions[self.round_number-1]
        interaction_number = Constants.interactions[self.round_number-1]

        for p in players: # set interaction number and round number
            p.interaction_number = interaction_number
            p.round_in_interaction = round_in_interaction
            p.treatment = (p.participant.id_in_session+p.interaction_number)%2
            p.group_size = len(p.group.get_players())

            if self.round_number == 1:
                p.participant.vars['returns'] = np.load('multiarmed_bandit/values.npy') - 15
                print(p.participant.vars['returns'])


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    choice = models.IntegerField(choices=[
        [0, 'A'],
        [1, 'B'],
        [2, 'C'],
        [3, 'D'],
        [4, 'E']]
    )
    mu = models.PositiveIntegerField()
    noise = models.FloatField()
    best_mu = models.PositiveIntegerField()
    ## it may be helpful to define treatment variable, making it easier for data analysis later

    timed_out = models.IntegerField()
    cum_payoff = models.CurrencyField()
    avg_payoff = models.FloatField()
    num_decisions = models.IntegerField()
    treatment = models.PositiveIntegerField()
    group_size = models.PositiveIntegerField()
    interaction_number = models.PositiveIntegerField()
    round_in_interaction = models.PositiveIntegerField()
    average_pay = models.FloatField() # average group payment
    max_pay = models.FloatField() # max group payment