from otree.api import *
c = cu

doc = 'Intro to experiment. Provides players brief introduction and gets informed consent. Then plays video instructions for first phase of experiment.'
class C(BaseConstants):
    # built-in constants
    NAME_IN_URL = 'Intro'
    PLAYERS_PER_GROUP = 2
    NUM_ROUNDS = 1
    # user-defined constants
    UW_TETON_BG_TEMPLATE = 'Intro/uw_teton_bg.html'
class Subsession(BaseSubsession):
    pass
def after_all_players_arrive(subsession: Subsession):
    session = subsession.session
    # Build a dict: code → list of player IDs
    buckets = {}
    for p in subsession.get_players():
        buckets.setdefault(p.couple_code.strip().lower(), []).append(p.id_in_subsession)
    # Now turn that into a list-of-lists matrix
    matrix = [ids for ids in buckets.values() if len(ids) == 2]
    # (You can handle odd or mistyped codes however you like here.)
    subsession.set_group_matrix(matrix)
    subsession.session.vars['couple_matrix'] = matrix
class Group(BaseGroup):
    pass
class Player(BasePlayer):
    couple_code = models.StringField(label='Please enter your couple code')
    my_field = models.IntegerField(choices=[])
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 Assignment(Page):
    form_model = 'player'
    form_fields = ['couple_code']
    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        group = player.group
        participant = player.participant
        # copy into participant.vars so it’s visible at grouping time
        player.participant.vars['couple_code'] = player.couple_code
class AssignmentWaitPage(WaitPage):
    wait_for_all_groups = True
    after_all_players_arrive = after_all_players_arrive
class Welcome(Page):
    form_model = 'player'
class Consent(Page):
    form_model = 'player'
# vars_for_template is deprecated in favor of python-based renderers
    @staticmethod
    def vars_for_template(player: Player):
        session = player.session
        subsession = player.subsession
        group = player.group
        participant = player.participant
        return dict(
            my_label=player.participant.label,
            my_group=player.group.id_in_subsession,
        )
class IntroEndPage(WaitPage):
    wait_for_all_groups = True
page_sequence = [Assignment, AssignmentWaitPage, Welcome, Consent, IntroEndPage]