from otree.api import *
import random

doc = """
The economic tasks  
"""

class C(BaseConstants):
    NAME_IN_URL = 'GroupDecisions'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 2
    INSTRUCTIONS_TEMPLATE = 'part3trustgames_PlayerAB/instructions.html'
    ENDOWMENT = cu(10)
    MAX_NOISE = 1
    BOTS_PER_GROUP = 1
    AGENTS_PER_GROUP = BOTS_PER_GROUP + 1
    MULTIPLIER = 3

class Subsession(BaseSubsession):
    pass

def creating_session(subsession: Subsession):
    for i, p in enumerate(subsession.get_players()):
        MyBot.create(player=p, agent_id=i)
        p.citizen_treatment = 'Yes'
        p.class_treatment = random.choice(\
                    ['$15,000-$29,999/year','$65,0000-$89,999/year'])
        p.age_treatment = '26-36'
        p.race_treatment = random.choice(\
                    ['White, non-Hispanic/Latino', 'African American or Black', 
                    'Hispanic/Latino', 'Asian']) 

class Group(BaseGroup):
    pass

class Player(BasePlayer):
       
    citizen_treatment  = models.StringField()
    age_treatment = models.StringField()
    race_treatment = models.StringField()
    class_treatment  = models.StringField()
    
    #Bot as Player A
    contrib = models.CurrencyField(
        min=0,
        max=30,
        doc="""Amount sent by P1""",
        label="",
    )
    agent_id = models.IntegerField(initial=1)
    bot_return = models.CurrencyField()
    points = models.CurrencyField()
    PA_send = models.CurrencyField()
    PA_tripled = models.CurrencyField()


    age = models.StringField(
       choices=[
        [1, '18-25'],
        [2, '26-36'],
        [3, '37-47'],
        [4, '48-59'],
        [5, '60 or above']],
        label= "Which age group do you belong to?"
)
    citizen = models.StringField(
        choices=[
            [1, 'Yes'],
            [2, 'No']],
        label= "Are you a US citizen?"
)
    race = models.StringField(
       choices=[
        [1, 'African American or Black'],
        [2, 'White, non-Hispanic/Latino'],
        [3, 'Hispanic/Latino'],
        [4, 'Asian'],
        [5, 'Other']],
    label= "What is your Ethnic/Racial background?"
)
    income = models.StringField(
       choices=[
        [1, '$0-$14,999/year'],
        [2, '$15,000-$29,999/year'],
        [3, '$30,000-$44,999/year'],
        [4, '$45,000–$64,999/year'], 
        [5, '$65,000-$89,999/year'],
        [6, '$90,000 or above']],
    label= "What is your income?"
)

#Defining Bot Behaviors - As Player A
class MyBot(ExtraModel):
    player = models.Link(Player)

    contrib = models.CurrencyField(
        min=0,
        doc="""Amount sent by P1""",
        label="Please enter an amount from 0 to 100:",
    )
    PA_send = models.CurrencyField()
    PA_tripled = models.CurrencyField()
    agent_id = models.IntegerField()

#Defining Bot Behaviors - As Player B    
def generate_contrib(max_contrib):
    import random
    
    def AB_diff(A_start, A_sent, B_sends_back, max_contrib):
        A = A_start - A_sent + B_sends_back
        B = max_contrib - B_sends_back
        return B - A

    A_start = C.ENDOWMENT
    contrib = 0
    A_sent = max_contrib / 3
    min_diff = abs(AB_diff(A_start, A_sent, contrib, max_contrib))
    for j in range(1, int(max_contrib+1)):
        new_diff = abs(AB_diff(A_start, A_sent, j, max_contrib))
        if new_diff < min_diff:
            contrib = j
            min_diff = new_diff
            
    return max(0, min(contrib, max_contrib))

# DEFINING PAYOUTS
    
def bot_return(player: Player):
    player.bot_return = MyBot.filter(player=player)[0].contrib
    return player.bot_return

def points(player: Player):
    if player.round_number==1:
        player.points = C.ENDOWMENT - player.contrib + player.bot_return
    else:
        player.points = player.PA_tripled - player.contrib
    player.payoff = float(player.points*100)*0.4
    return player.points

def set_payoffs(player: Player):
    bot = MyBot.filter(player=player)[0]
    max_contrib = player.contrib * C.MULTIPLIER
    bot.contrib = generate_contrib(max_contrib)

def set_payoffs_bot(player: Player):
    bot = MyBot.filter(player=player)[0]
    bot.PA_send = random.randint(0,10)
    bot.PA_tripled = bot.PA_send*C.MULTIPLIER
    player.PA_send = bot.PA_send
    player.PA_tripled = bot.PA_tripled
    

class Survey(Page):
     form_model = 'player'
     form_fields = ['age', 'citizen', 'race','income'] 
     
     
     @staticmethod
     def is_displayed(player:Player):
         return player.round_number == 1

class WaitForPlayer_Survey(Page):
        
   @staticmethod
   def get_timeout_seconds(player: Player):
        return random.randint(5, 12) 

class WaitForPlayerA_Task(Page):
    
  @staticmethod
  def is_displayed(player:Player):
      return player.round_number == 2
  
  @staticmethod
  def get_timeout_seconds(player: Player):
      return random.randint(6, 20)     

class Send(Page):

    form_model = 'player'
    form_fields = ['contrib']
        
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        if player.round_number==1:
            set_payoffs(player)
 
    @staticmethod
    def vars_for_template(player: Player):
        if (player.round_number == 2) & (player.field_maybe_none('PA_send')==None):
            set_payoffs_bot(player)
            
    @staticmethod
    def error_message(player: Player, values):
        if player.round_number==1:
            if values['contrib'] > 10:
                return "You cannot send more than 10 points."
        else:
            if values['contrib'] > player.PA_tripled:
                  return "You cannot send more than " + str(player.PA_tripled) + "."


class WaitForPlayer_Task(Page):
    
  @staticmethod
  def is_displayed(player:Player):
      return player.round_number == 1
  
  @staticmethod
  def get_timeout_seconds(player: Player):
      return random.randint(6, 14)
  
class Results(Page):

    @staticmethod
    def vars_for_template(player: Player):
        prev_player = player.in_round(1)
        bot = MyBot.filter(player=player)[0]
        bot_return(prev_player)
        points(player)        
        total_payoff = prev_player.points + player.points
        dollar_paid = round(float(total_payoff)*0.4, 2)
        return dict(prev_player= prev_player,
                    bot=bot,
                    total_payoff=total_payoff,
                    dollar_paid=dollar_paid)

page_sequence = [Survey,WaitForPlayer_Survey,WaitForPlayerA_Task, Send,
                 WaitForPlayer_Task, Results]