from otree.api import *
c = cu

doc = 'Real-effort task (RET) portion of the experiment. RET is counting the number of 1s in a matrix of 1s and 0s.'
class C(BaseConstants):
    # built-in constants
    NAME_IN_URL = 'RET_1'
    PLAYERS_PER_GROUP = 2
    NUM_ROUNDS = 20
    # user-defined constants
    TIMER_TEXT = 'Time left to complete the matrix task:'
    UW_TETON_BG_TEMPLATE = 'RET_1/uw_teton_bg.html'
class Subsession(BaseSubsession):
    pass
def after_all_players_arrive(subsession: Subsession):
    session = subsession.session
    # pull the matrix you stored in the code‐entry app
    matrix = subsession.session.vars.get('couple_matrix')
    if not matrix:
        # safety check in case someone forgot to run grouping first
        raise RuntimeError("No couple_matrix in session.vars")
    # re‐apply it in this app's subsession
    subsession.set_group_matrix(matrix)
class Group(BaseGroup):
    budget = models.IntegerField()
    correct = models.IntegerField(initial=0)
    earner_id = models.FloatField()
def check_func(group: Group):
    session = group.session
    # Access all players in the current group
    players = group.get_players()
    
    # Exit check for conditions that don't require individual updates
    if group.round_number > 1 and group.correct == session.config['num_correct']:
        # Directly carry over the correct count from the previous round
        group.correct = group.in_round(group.round_number - 1).correct
        return
    
    # Identify the player who did the task in the first round
    ptask = group.get_player_by_id(group.in_round(1).earner_id)
    
    # Initialize or update correct counts for players
    for p in players:
        # Ensure that player_sum is not None, treat None as an incorrect guess
        if p.player_sum is None:
            p.last_guess_correct = False
            current_guess_correct = 0  # This will treat None as an incorrect guess
        else:
            # Compare player_sum with matrix_sum and update correct and feedback accordingly
            current_guess_correct = int(p.player_sum == p.matrix_sum)
            p.last_guess_correct = (p.player_sum == p.matrix_sum)
    
        # Update the correct count based on the current guess
        if group.round_number == 1:
            p.correct = current_guess_correct
        else:
            previous_correct = p.in_round(group.round_number - 1).correct
            p.correct = previous_correct + current_guess_correct
    
        # Update the group.correct counter if this is the player who did the task
        if p == ptask:
            group.correct = p.correct
def RET_payoff(group: Group):
    session = group.session
    # Defining player that did the task
    ptask = group.get_player_by_id(group.in_round(1).earner_id)
    
    # Setting the budget based on how many were correct and the pay per correct answer
    group.budget = ptask.correct * session.config['x_rate']
    
    # Defining the players
    players = group.get_players()
    
    # Store budget at the participant level for the next app
    for p in players:
        p.participant.budget_1 = group.budget
def assign_earner_1(group: Group):
    import random
    
    # Randomly choosing one of the group members to be the earner in round 1
    r = random.randint(1,2)
    
    # Gather players 
    players = group.get_players()
    
    # Initialize earner_id at the group level
    group.earner_id = None
    
    # Assign roles based on random draw
    for p in players:
        if p.id_in_group == r:
            p.participant.earner_round_1 = 1
            group.earner_id = p.id_in_group  # Store the earner's id_in_group at the group level
        else:
            p.participant.earner_round_1 = 0
    
        p.exit_1 = 0
def exit(group: Group):
    # Defining players
    players = group.get_players()
    
    # Setting the exit condition to avoid duplicating the results page
    for p in players:
        p.exit_1 = 1
class Player(BasePlayer):
    matrix_sum = models.IntegerField(initial=1)
    player_sum = models.IntegerField(initial=0, label='How many 1s are in the above matrix?', min=0)
    correct = models.IntegerField(initial=0, min=0)
    last_guess_correct = models.BooleanField()
    exit_1 = models.FloatField(initial=0)
    cq = models.IntegerField(choices=[[1, 'The money goes to the participant who earned it through the matrix task'], [2, 'The money is split evenly between both partners'], [3, 'The participant chosen to spend money keeps all remaining funds']], label='What happens to the money in the joint account at the end of the study?', widget=widgets.RadioSelect)
    cq_try2 = models.IntegerField(choices=[[1, 'The money goes to the participant who earned it through the matrix task'], [2, 'The money is split evenly between both partners'], [3, 'The participant chosen to spend money keeps all remaining funds']], label='What happens to the money in the joint account at the end of the study?', widget=widgets.RadioSelect)
def player_sum_max(player: Player):
    session = player.session
    return session.config['m_size'] * session.config['m_size']
def custom_export(players):
    # Header
    yield [
        'session_code',
        'participant_id_in_session',
        'participant_code',
        'id_in_group',
        'group_number',
        'earner_round_1',
        'budget_1'
    ]
    
    # Rows
    for p in players:
        pp = p.participant
        pg = p.group
        ps = p.session
    
        # Pull participant fields
        earner_round_1 = p.participant.vars.get('earner_round_1', '')
        budget_1 = p.participant.vars.get('budget_1', '')
    
        yield [
            ps.code,
            pp.id_in_session,
            pp.code,
            p.id_in_group,
            pg.id_in_subsession,
            earner_round_1,
            budget_1
        ]
class GroupingWaitPage(WaitPage):
    wait_for_all_groups = True
    after_all_players_arrive = after_all_players_arrive
class RoleAssignmentWaitPage(WaitPage):
    after_all_players_arrive = assign_earner_1
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        if group.round_number == 1: 
            return True
class RoleAssignment(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        if group.round_number == 1: 
            return True
# vars_for_template is deprecated in favor of python-based renderers
    @staticmethod
    def vars_for_template(player: Player):
        session = player.session
        # Calculate the maximum possible earnings from the task
        max_budget = session.config['x_rate'] * session.config['num_correct']
        
        # Passing Python variables to the HTML template of the page
        return dict(max_budget = max_budget)
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        session = player.session
        participant = player.participant
        # Set a global 'task_min' minute timer for all the task pages
        participant = player.participant
        import time
        
        # Defining the global timer amount using the session.config
        participant.expiry = time.time() + session.config['task_min']*60
class Task(Page):
    timer_text = 'Time left to complete the matrix task:'
    form_model = 'player'
    form_fields = ['player_sum']
    @staticmethod
    def is_displayed(player: Player):
        session = player.session
        group = player.group
        participant = player.participant
        import time
        
        # Display if participant is the earner AND it's round 1 AND global timer isn't 0
        if player.participant.earner_round_1 == 1 and group.round_number == 1 and max(participant.expiry - time.time(), 0) > 3:
            return True
        
        # Display if participant is the earner AND it's after round 1 AND is not done with the task AND global timer isn't 0
        elif player.participant.earner_round_1 == 1 and group.in_round(group.round_number - 1).correct != session.config['num_correct'] and max(participant.expiry - time.time(), 0) > 3:
            return True
# vars_for_template is deprecated in favor of python-based renderers
    @staticmethod
    def vars_for_template(player: Player):
        session = player.session
        import random
        
        # Generate session.config.m_size x session.config.m_size matrix with random 0s and 1s
        matrix = [[random.randint(0, 1) for _ in range(session.config['m_size'])] for _ in range(session.config['m_size'])]
        
        # Calculate the sum of all 1s in the matrix
        matrix_sum = sum(sum(row) for row in matrix)
        
        # Store the sum in the player variable
        player.matrix_sum = matrix_sum
        
        # Passing Python variables to the HTML template of the page
        return dict(
            matrix=matrix,
            matrix_sum=matrix_sum,
        )
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        if timeout_happened:
            # Default value for any form fields if timeout occurs, because otherwise they may be left null
            player.player_sum = 0
    @staticmethod
    def get_timeout_seconds(player: Player):
        participant = player.participant
        # Using the global timer set at RoleAssignment
        participant = player.participant
        import time
        
        # Passing the global timer to the HTML template of the page
        return participant.expiry - time.time()
class EarningWaitPage(WaitPage):
    after_all_players_arrive = check_func
    @staticmethod
    def is_displayed(player: Player):
        participant = player.participant
        # Display if earner
        if participant.earner_round_1 == 1:
            return True
class Feedback(Page):
    timer_text = 'Time left to complete the matrix task:'
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        session = player.session
        group = player.group
        participant = player.participant
        import time
        
        # Display if participant is the earner AND it's round 1 AND global timer isn't 0
        if participant.earner_round_1 == 1 and group.round_number == 1 and max(participant.expiry - time.time(), 0) > 3:
            return True
        
        # Display if participant is the earner AND it's after round 1 AND is not done with the task AND global timer isn't 0
        elif participant.earner_round_1 == 1 and group.in_round(group.round_number - 1).correct != session.config['num_correct'] and max(participant.expiry - time.time(), 0) > 3:
            return True
    @staticmethod
    def get_timeout_seconds(player: Player):
        participant = player.participant
        # Using the global timer set at RoleAssignment
        participant = player.participant
        import time
        
        # Passing the global timer to the HTML template of the page
        return participant.expiry - time.time()
class TimeExpired(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        participant = player.participant
        import time
        
        # Only display this page if the participant ran out of time on the matrix task
        return (player.participant.expiry - time.time()) <= 3
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        session = player.session
        group = player.group
        # Used to manually reach the exit condition to get to the results screen if timed out
        group.round_number = session.config['num_rounds']
class NotEarningWaitPage(WaitPage):
    title_text = 'Please wait for further instructions.'
    @staticmethod
    def is_displayed(player: Player):
        participant = player.participant
        # If the participant is not the earner, have them sit on this wait page
        if participant.earner_round_1 != 1:
            return True
class PayoffWaitPage(WaitPage):
    after_all_players_arrive = RET_payoff
    @staticmethod
    def is_displayed(player: Player):
        session = player.session
        group = player.group
        # Only reach this page if the group has gotten enough matrices correct OR they hit the max number of rounds
        if group.correct == session.config['num_correct'] or group.round_number == session.config['num_rounds']:
            return True
class ExitWaitPage(WaitPage):
    wait_for_all_groups = True
    @staticmethod
    def is_displayed(player: Player):
        session = player.session
        group = player.group
        # Display if they are done or we hit the round limit
        if group.correct == session.config['num_correct'] or group.round_number == session.config['num_rounds']:
            return True
        else:
            return False
class Results(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        session = player.session
        group = player.group
        participant = player.participant
        # Display if they are done or we hit the round limit
        if group.correct == session.config['num_correct'] or group.round_number == session.config['num_rounds']:
            return True
        else:
            return False
        
        # Only show the page if we haven't already shown it
        if player.participant.vars.get('my_page_shown', False):
            return False
        else:
            player.participant.vars['my_page_shown'] = True
            return True
class JointAccount(Page):
    form_model = 'player'
    form_fields = ['cq']
    @staticmethod
    def is_displayed(player: Player):
        session = player.session
        group = player.group
        # Display if they are done or we hit the round limit
        if group.correct == session.config['num_correct'] or group.round_number == session.config['num_rounds']:
            return True
        else:
            return False
class CQ(Page):
    form_model = 'player'
    form_fields = ['cq_try2']
    @staticmethod
    def is_displayed(player: Player):
        session = player.session
        group = player.group
        participant = player.participant
        # Display if they are done or we hit the round limit and If participant got comprehension question wrong, show them this page
        if group.correct == session.config['num_correct'] and player.cq != 2 or group.round_number == session.config['num_rounds'] and player.cq != 2:
            return True
        else:
            return False
class JointAccountEnd(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        session = player.session
        group = player.group
        # Display if they are done or we hit the round limit
        if group.correct == session.config['num_correct'] or group.round_number == session.config['num_rounds']:
            return True
        else:
            return False
    @staticmethod
    def app_after_this_page(player: Player, upcoming_apps):
        # Used to skip directly to the next app
        return "Auction_instructions"
page_sequence = [GroupingWaitPage, RoleAssignmentWaitPage, RoleAssignment, Task, EarningWaitPage, Feedback, TimeExpired, NotEarningWaitPage, PayoffWaitPage, ExitWaitPage, Results, JointAccount, CQ, JointAccountEnd]