from otree.api import *
import random

doc = """
This is a standard 2-player trust game where the amount sent by player 1 gets
tripled. The trust game was first proposed by......
"""


class C(BaseConstants):
    NAME_IN_URL = 'Games'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 2
    INSTRUCTIONS_TEMPLATE = 'part2diss_trustgames/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

# treatments for players and bots
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):

#Trust allocations    
    citizen_treatment  = models.StringField()
    age_treatment = models.StringField()
    race_treatment = models.StringField()
    class_treatment  = models.StringField()
    
    contrib = models.CurrencyField(
        min=0,
        max=C.ENDOWMENT,
        doc="""Amount sent by P1""",
        label="Please enter an amount from 0 to 10 you would like to send Player B:",
    )
    agent_id = models.IntegerField(initial=1)
    bot_return = models.CurrencyField()
    points = models.CurrencyField()

# Survey 
    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 Behavior 

class MyBot(ExtraModel):
    player = models.Link(Player)

    # these fields should match what's defined on the Player, so that you can
    # loop over player and bot instances interchangeably.
    contrib = models.CurrencyField(
        min=0,
        doc="""Amount sent by P1""",
        label="Please enter an amount from 0 to 100:",
    )
    agent_id = models.IntegerField()
    
def generate_contrib(max_contrib):
    import random

    #contrib = (max_contrib/2) + random.randint(-C.MAX_NOISE, C.MAX_NOISE)
    
    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
    
    # constrain it between 0 and endowment
    #return max(0, min(contrib, player.contrib * MULTIPLIER))
    return max(0, min(contrib, max_contrib))

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

def points(player: Player):
    player.points = C.ENDOWMENT - player.contrib + player.bot_return
    return player.points

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


class WaitForSurvey1(Page):
    
    """
    This is just for show, to make it feel more realistic.
    Also, note it's a Page, not a WaitPage.
    Removing this page won't affect functionality.
    """
    
    @staticmethod
    def is_displayed(player:Player):
        return player.round_number == 1
        
    @staticmethod
    def get_timeout_seconds(player: Player):
        return random.randint(3, 5) 

class Send1(Page):
    """This page is only for P1
    P1 sends amount (all, some, or none) to P2
    This amount is tripled by experimenter,
    i.e if sent amount by P1 is 5, amount received by P2 is 15"""

    form_model = 'player'
    form_fields = ['contrib']

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        set_payoffs(player)

class WaitForBots(Page):
    """
    This is just for show, to make it feel more realistic.
    Also, note it's a Page, not a WaitPage.
    Removing this page won't affect functionality.
    """

    @staticmethod
    def get_timeout_seconds(player: Player):
        return random.randint(3, 5) 

class Results(Page):
    @staticmethod
    def vars_for_template(player: Player):
        prev_player = player.in_round(1)
        bot = MyBot.filter(player=player)[0]
        br = bot_return(player)
        pts = 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,
                    bot_return=br,
                    total_payoff=total_payoff,
                    dollar_paid=dollar_paid)

page_sequence = [Survey, WaitForSurvey1, Send1, WaitForBots, Results]