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_practice'
    PLAYERS_PER_GROUP = 2
    NUM_ROUNDS = 3
    # user-defined constants
    NUM_CORRECT = 3
    UW_TETON_BG_TEMPLATE = 'RET_practice/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):
    # 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 == C.NUM_CORRECT:
        # Directly carry over the correct count from the previous round
        group.correct = group.in_round(group.round_number - 1).correct
        return
    
    # 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
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)
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'
    ]
    
    # Rows
    for p in players:
        pp = p.participant
        pg = p.group
        ps = p.session
    
        yield [
            ps.code,
            pp.id_in_session,
            pp.code,
            p.id_in_group,
            pg.id_in_subsession
        ]
class GroupingWaitPage(WaitPage):
    wait_for_all_groups = True
    after_all_players_arrive = after_all_players_arrive
class Intro(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        if group.round_number == 1: 
            return True
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        pass
class Task(Page):
    timeout_seconds = 60
    form_model = 'player'
    form_fields = ['player_sum']
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        # Display if it's round 1
        if group.round_number == 1:
            return True
        
        # Display if it's after round 1 AND is not done with the task
        elif group.in_round(group.round_number - 1).correct != C.NUM_CORRECT:
            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 C.M_SIZExC.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
class EarningWaitPage(WaitPage):
    after_all_players_arrive = check_func
class Feedback(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        participant = player.participant
        # Display if participant is the earner AND it's round 1
        if group.round_number == 1:
            return True
        
        # Display if participant is the earner AND it's after round 1 AND is not done with the task
        elif group.in_round(group.round_number - 1).correct != C.NUM_CORRECT:
            return True
        
        # Don't display if final round
        if group.round_number == C.NUM_CORRECT:
            return False
class Results(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        # Display if they are done or we hit the round limit
        if group.round_number == C.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 "RET_1"
class RET_practiceEndPage(WaitPage):
    wait_for_all_groups = True
    @staticmethod
    def is_displayed(player: Player):
        group = player.group
        # Display if they are done or we hit the round limit
        if group.round_number == C.NUM_ROUNDS:
            return True
        else:
            return False
page_sequence = [GroupingWaitPage, Intro, Task, EarningWaitPage, Feedback, Results, RET_practiceEndPage]