from otree.api import *
import random

doc = """
Asynchronous 2-player sequential game (players can play at different times),
where we guarantee players never have to wait.
The example used is an ultimatum game.  
The first player who arrives is assigned to be P1.
If the next player arrives after P1 has made a decision, he will be paired with P1 and see P1's decision.
Otherwise, he will be P1 in a new group.
Worst-case scenario is that all players arrive around the same time, 
and therefore everyone gets assigned to be P1.

This game doesn't use oTree groups. 
Rather, it stores the partner's ID in a player field.
"""

class C(BaseConstants):
    NAME_IN_URL = 'AI_ult'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 1
    ENDOWMENT = cu(79)
    TREATMENTS = ['ai', 'human']

class Subsession(BaseSubsession):
    pass

def creating_session(subsession: Subsession):
    players = subsession.get_players()
    num_players = len(players)

    # Alternate between the two treatments for each player
    for i, player in enumerate(players, start=1):
        if i % 2 == 0:
            player.condition = 'human'
        else:
            player.condition = 'ai'
        print(f"Assigning player {player.id_in_group} to condition {player.condition}")

class Group(BaseGroup):
    pass

class Player(BasePlayer):
    condition = models.StringField()
    is_p1 = models.BooleanField()
    offer = models.CurrencyField(
        min=0,
        max=C.ENDOWMENT,
        label="What is <b>your offer</b> to the AI bot? (Amount between $0 and $79)",
        doc="This field is only used if the player is P1",
    )
    potsize = models.CurrencyField(
        min=5,
        max=90,
        label="How much will you tell the AI bot the <b>pot size</b> is? (Amount between $5-$90)",
        doc="This field is only used if the player is P1",
    )

    partner_id = models.IntegerField(
        doc="This field is only used if the player is P2. It stores the ID of P1",
    )
    accepted = models.BooleanField(
        label="Do you accept Player 1's offer?", doc="This field is only used if the player is P2",
    )
    bot_accepted = models.BooleanField()
    quiz1 = models.IntegerField(label='The size of the pot is $____.')
    
    quiz2a = models.StringField(
    label="The AI bot is told that the pot size ranges between _____.",
    choices=['$5 - $100', '$0 - $70', '$5 - $90'],)

    quiz2h = models.StringField(
    label="The other participant is told that the pot size ranges between _____.",
    choices=['$5 - $100', '$0 - $70', '$5 - $90'],)

    quiz3 = models.StringField(
    label="As the first participant, I will receive ____.",
    choices=['All of the pot money', 'The pot money minus my offer amount', 'None of the pot money'],)
    
    quiz4a = models.BooleanField(label="The AI bot knows the actual size of the pot.")
    quiz4h = models.BooleanField(label="The other participant knows the actual size of the pot.")
       
def set_payoff(player: Player):
    if player.bot_accepted:
        player.payoff = C.ENDOWMENT - player.offer
    else:
        player.payoff = cu(0)

class AIU_Intro(Page):
    form_model = 'player'
    # form_fields = ['quiz1', 'quiz2a', 'quiz3', 'quiz4a']
    
    @staticmethod
    def is_displayed(player):
        return player.condition == 'ai'
        
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        session = player.session

class AIU_Instruct(Page):
    form_model = 'player'
    form_fields = ['quiz1', 'quiz2a', 'quiz3', 'quiz4a']
    
    @staticmethod
    def is_displayed(player):
        return player.condition == 'ai'
        
    @staticmethod
    def error_message(player: Player, values):
        solutions = dict(quiz1=C.ENDOWMENT, quiz2a='$5 - $90', quiz3='The pot money minus my offer amount', quiz4a=False)

        if values != solutions:
            return "One or more answers were incorrect. Please reread the instructions and try again."
    
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        session = player.session

class AIU_Offer(Page):
    form_model = 'player'
    form_fields = ['offer', 'potsize']

    @staticmethod
    def is_displayed(player):
        return player.condition == 'ai'
        
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        session = player.session
        
        # Rest of your code for the offer...
        half_pot_size = player.potsize / 2
        half_endowment = (C.ENDOWMENT - 1) / 2

        # Condition 1: Accept if offer is equal to or greater than one less than half of the endowment
        if player.offer >= half_endowment:
            player.bot_accepted = True

        # Condition 2: Reject if offer is less than half of the stated pot size
        elif player.offer < half_pot_size:
            player.bot_accepted = False

        # Condition 3: Randomly accept or reject if in between
        else:
            player.bot_accepted = random.choice([True, False])

        # Ensure the payoff is calculated based on this decision
        set_payoff(player)

class AIU_Results(Page):
    @staticmethod
    def is_displayed(player):
        return player.condition == 'ai'
        
    @staticmethod
    def vars_for_template(player: Player):
        bot_decision = "Accepted" if player.bot_accepted else "Rejected"
        bot_payoff = player.offer if player.bot_accepted else cu(0)
        player_payoff=player.payoff
        return dict(
            bot_decision=bot_decision,
            bot_payoff=bot_payoff
        )
    
class HumanU_Intro(Page):
    form_model = 'player'
    # form_fields = ['quiz1', 'quiz2h', 'quiz3', 'quiz4h']
    
    @staticmethod
    def is_displayed(player):
        return player.condition == 'human'
    
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        session = player.session
    
class HumanU_Instruct(Page):
    form_model = 'player'
    form_fields = ['quiz1', 'quiz2h', 'quiz3', 'quiz4h']
    
    @staticmethod
    def is_displayed(player):
        return player.condition == 'human'
    
    @staticmethod
    def error_message(player: Player, values):
        solutions = dict(quiz1=C.ENDOWMENT, quiz2h='$5 - $90', quiz3='The pot money minus my offer amount', quiz4h=False)
        
        if values != solutions:
            return "One or more answers were incorrect. Please reread the instructions and try again."
        
class HumanU_Offer(Page):
    form_model = 'player'
    form_fields = ['offer', 'potsize']

    @staticmethod
    def is_displayed(player):
        return player.condition == 'human'
    
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        session = player.session
        half_pot_size = player.potsize / 2
        half_endowment = (C.ENDOWMENT - 1) / 2
    
        # Condition 1: Accept if offer is equal to or greater than one less than half of the endowment
        if player.offer >= half_endowment:
            player.bot_accepted = True

        # Condition 2: Reject if offer is less than half of the stated pot size
        elif player.offer < half_pot_size:
            player.bot_accepted = False

        # Condition 3: Randomly accept or reject if in between
        else:
            player.bot_accepted = random.choice([True, False])

        # Ensure the payoff is calculated based on this decision
        set_payoff(player)

  
class HumanU_Results(Page):
    @staticmethod
    def is_displayed(player):
        return player.condition == 'human'
    
    @staticmethod
    def vars_for_template(player: Player):
        bot_decision = "Accepted" if player.bot_accepted else "Rejected"
        bot_payoff = player.offer if player.bot_accepted else cu(0)
        player_payoff=player.payoff
        return dict(
            bot_decision=bot_decision,
            bot_payoff=bot_payoff
        )
    
page_sequence = [AIU_Intro, HumanU_Intro, AIU_Instruct, HumanU_Instruct, AIU_Offer, HumanU_Offer, AIU_Results, HumanU_Results]