from otree.api import *
import random

doc = """
Ultimatum game
AI vs. Human
Fading mediation
NO Moderation

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

class C(BaseConstants):
    NAME_IN_URL = 'Ultimatum'
    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()
    last_condition = None

    for player in players:
        # Check if the previous player was a dropout
        prev_player = player.in_all_rounds()[-2] if len(player.in_all_rounds()) > 1 else None
        if prev_player and prev_player.is_dropout:
            # Assign this player the same condition as the previous dropout player
            player.condition = prev_player.condition
        else:
            # Alternate between 'ai' and 'human' conditions
            if last_condition == 'ai':
                player.condition = 'human'
            else:
                player.condition = 'ai'
        print(f"Assigning player {player.id_in_group} to condition {player.condition}")
        
        # Update last_condition
        last_condition = player.condition

   # print(f"Assigning player {player.id_in_group} to condition {player.condition}")

    # num_players = len(players)
    # last_player = players[-1] if players else None

    # for i, player in enumerate(players, start=1):
    #     if last_player and last_player.is_dropout:
    #         # Assign this player as a makeup for the last dropout
    #         player.condition = last_player.condition
    #     else:
    #     # Alternate between the two treatments for each player
    #         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,
        doc="This field is only used if the player is P1",
    )
    potsize = models.CurrencyField(
        min=5,
        max=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.")

    is_dropout = models.BooleanField(initial=False)
       
def set_payoff(player: Player):
    if player.bot_accepted:
        player.payoff = C.ENDOWMENT - player.offer
    else:
        player.payoff = cu(0)

### AI PARTNERSHIP VERSION ###    
class U1Intro(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 U1Instruct(Page):
    timeout_seconds = 180

    form_model = 'player'
    form_fields = ['quiz1', 'quiz2a', 'quiz3', 'quiz4a']

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        # If timeout happens, mark as dropout
        if timeout_happened:
            player.is_dropout = True
        return
    
    @staticmethod
    def is_displayed(player):
        return player.condition == 'ai' and not player.is_dropout
        
    @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."
    
class U1Offer(Page):
    timeout_seconds = 180
    
    form_model = 'player'
    form_fields = ['offer', 'potsize']

    @staticmethod
    def is_displayed(player):
        return player.condition == 'ai' and not player.is_dropout

    @staticmethod
    def vars_for_template(player: Player):
        offer_label = "What is <b>your offer</b> to the AI bot? (Amount between $0 and $79)"
        potsize_label = "How much will you tell the AI bot the <b>pot size</b> is? (Amount between $5-$90)"
        return {
            'offer_label': offer_label,
            'potsize_label': potsize_label
        }

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        # If timeout happens, mark as dropout
        if timeout_happened:
            player.is_dropout = True
            return
        session = player.session
        
        # Rest of your code for the offer...
        half_pot_size = player.potsize / 2
        half_endowment = (C.ENDOWMENT - 1) / 2

        # Condition 0: Accept if offer is equal to or greater than 1
        if player.offer >= 1:
            player.bot_accepted = True

        # Condition 1: Accept if offer is equal to or greater than one less than half of the endowment
        elif 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 U1Results(Page):
    @staticmethod
    def is_displayed(player):
        return player.condition == 'ai' and not player.is_dropout
          
    @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
        )

### HUMAN PARTNERSHIP VERSION ###    
class U2Intro(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 U2Instruct(Page):
    timeout_seconds = 180
    
    form_model = 'player'
    form_fields = ['quiz1', 'quiz2h', 'quiz3', 'quiz4h']
    
    @staticmethod
    def is_displayed(player):
        return player.condition == 'human' and not player.is_dropout
    
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        # If timeout happens, mark as dropout
        if timeout_happened:
            player.is_dropout = True
        return
    
    @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 U2Offer(Page):
    timeout_seconds = 180
    form_model = 'player'
    form_fields = ['offer', 'potsize']

    @staticmethod
    def is_displayed(player):
        return player.condition == 'human' and not player.is_dropout
    
    @staticmethod
    def vars_for_template(player: Player):
        offer_label = "What is <b>your offer</b> to the other participant? (Amount between $0 and $79)"
        potsize_label = "How much will you tell the other participant the <b>pot size</b> is? (Amount between $5-$90)"
        return {
            'offer_label': offer_label,
            'potsize_label': potsize_label
        }

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        if timeout_happened:
            player.is_dropout = True
            return

        session = player.session
        half_pot_size = player.potsize / 2
        half_endowment = (C.ENDOWMENT - 1) / 2
    
        # Condition 0: Accept if offer is equal to or greater than 1
        # HARD CODED TO ACCEPT EVERYTHING!
        if player.offer >= 1:
            player.bot_accepted = True

        # Condition 1: Accept if offer is equal to or greater than one less than half of the endowment
        elif 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 U2Results(Page):
    @staticmethod
    def is_displayed(player):
        return player.condition == 'human' and not player.is_dropout

    @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 UTimeout(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.is_dropout

    @staticmethod
    def error_message(player: Player, values):
        return "Cannot proceed past this page"
    
page_sequence = [U1Intro, U2Intro, U1Instruct, U2Instruct, U1Offer, U2Offer, UTimeout, U1Results, U2Results]