from otree.api import *
import random
import string


doc = """
Consent -> Instructions -> Comprehension -> Message
-> (LoadTask if load) -> Demographics
-> (SwitchMessage if switch) -> Decision
-> (Recall if load) -> Beliefs -> ThankYou

Load: alternating by participant (even id_in_subsession = load).
Switch: deterministic (first half no-switch, second half switch).
"""


class C(BaseConstants):
    NAME_IN_URL = 'myapp'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 1
    PASSWORD_LEN = 7


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # Consent
    consent = models.BooleanField(
        label='I am at least 18 years old and I consent to participate in this study.'
    )

    # Comprehension questions
    cq1 = models.StringField(
        label="1) Are messages binding in this experiment?",
        choices=[['binding', 'Yes, messages are binding'],
                 ['not_binding', 'No, messages are not binding']],
        widget=widgets.RadioSelect
    )
    cq2 = models.StringField(
        label="2) What payoff would you receive if Player A chooses OUT and you choose DONT ROLL?",
        choices=[['zero', 'You will receive 5 points.'],
                 ['fourteen', 'You will receive 14 points.']],
        widget=widgets.RadioSelect
    )
    cq3 = models.StringField(
        label="3) Player A will certainly know whether a switch occurred or not.",
        choices=[['yes', 'True'], ['no', 'False']],
        widget=widgets.RadioSelect
    )
    cq4 = models.StringField(
        label="4) Which choice gives you the highest guaranteed payoff if Player A chooses IN?",
        choices=[['roll', 'ROLL'], ['dont_roll', "DON’T ROLL"]],
        widget=widgets.RadioSelect
    )

    # Message
    message = models.LongStringField(label="Please write your message to Player A:")

    # Treatments
    load_condition = models.IntegerField(initial=0)     # set in creating_session()
    password = models.StringField(initial='', blank=True)
    password_recall = models.StringField(
        label="Please type the 7-character string you saw earlier (best effort):",
        initial='',
        blank=True
    )
    memory_strategy = models.StringField(
        label="When trying to remember the 7-character string, which of the following best describes what you did?",
        choices=[
            ['memory_only', 'I relied only on my memory.'],
            ['mental_repetition', 'I repeated the string to myself mentally.'],
            ['wrote_down', 'I wrote the string down.'],
            ['other_aid', 'I used another aid (e.g. note, phone, screenshot, or similar).'],
            ['prefer_not', 'Prefer not to say.'],
        ],
        widget=widgets.RadioSelect
    )

    switch_condition = models.IntegerField(initial=0)   # set in creating_session()
    switched_message = models.LongStringField(initial='', blank=True)

    # Demographics
    age = models.IntegerField(label="Age:", min=18, max=120, blank=True)
    gender = models.StringField(
        label="Gender:",
        choices=[['female', 'Female'],
                 ['male', 'Male'],
                 ['nonbinary', 'Non-binary / third gender'],
                 ['prefer_not', 'Prefer not to say']],
        blank=True
    )
    education = models.StringField(
        label="Highest completed education:",
        choices=[['hs', 'High school or equivalent'],
                 ['ba', "Bachelor's degree"],
                 ['ma', "Master's degree"],
                 ['phd', "PhD / doctorate"],
                 ['other', 'Other'],
                 ['prefer_not', 'Prefer not to say']],
        blank=True
    )

    # Decision
    roll_choice = models.StringField(
        label="Please choose:",
        choices=[['roll', 'ROLL'], ['dont_roll', "DON’T ROLL"]],
        widget=widgets.RadioSelect
    )

    # Second-order beliefs (Vanberg-style 5-point categorical scale)
    belief_roll_cat = models.StringField(
        label="Please select only one option. Player A will expect that I...",
        choices=[
            ['c_roll', 'Certainly choose ROLL'],
            ['p_roll', 'Probably choose ROLL'],
            ['unsure', 'Unsure'],
            ['p_dont', "Probably choose DON’T ROLL"],
            ['c_dont', "Certainly choose DON’T ROLL"],
        ],
        widget=widgets.RadioSelect
    )

    # Email field
    outcome_email = models.StringField(blank=True, label="Email address (optional)")


# ✅ IMPORTANT: in your oTree version, this MUST be a module-level function
def creating_session(subsession: Subsession):
    players = subsession.get_players()
    n = len(players)

    # use ceil(n/2) so with odd n you still have a non-empty first half
    half = (n + 1) // 2

    for p in players:
        # deterministic switch
        p.switch_condition = 0 if p.id_in_subsession <= half else 1

        # alternating load across ALL participants: even id => load
        p.load_condition = 1 if (p.id_in_subsession % 2 == 0) else 0

        # pre-generate password for load participants
        if p.load_condition == 1:
            chars = string.ascii_uppercase + string.digits
            p.password = ''.join(random.choice(chars) for _ in range(C.PASSWORD_LEN))
        else:
            p.password = ''


# ---------- helpers ----------

def _ensure_message_pool(session):
    if 'message_pool' not in session.vars:
        session.vars['message_pool'] = []


def _add_message_to_pool(player: Player):
    _ensure_message_pool(player.session)
    player.session.vars['message_pool'].append({'pid': player.id_in_subsession, 'msg': player.message})


def _pick_message_for_switch_player(player: Player):
    _ensure_message_pool(player.session)
    pool = player.session.vars['message_pool']

    # eligible messages from first half (no-switch participants)
    n = player.session.num_participants
    half = (n + 1) // 2
    eligible = [x for x in pool if x.get('pid', 999999) <= half]

    if eligible:
        choice = random.choice(eligible)
        player.switched_message = choice.get('msg', '') or ''
    else:
        player.switched_message = "No message is available yet. Please continue."


# ---------- PAGES ----------

class MyPage(Page):
    form_model = 'player'
    form_fields = ['consent']

    @staticmethod
    def error_message(player: Player, values):
        if not values.get('consent'):
            return "You must provide consent to participate in this study."


class Instructions(Page):
    pass


class Comprehension(Page):
    form_model = 'player'
    form_fields = ['cq1', 'cq2', 'cq3', 'cq4']


class Message(Page):
    form_model = 'player'
    form_fields = ['message']

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        _add_message_to_pool(player)


class LoadTask(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.load_condition == 1


class Demographics(Page):
    form_model = 'player'
    form_fields = ['age', 'gender', 'education']

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        if player.switch_condition == 1:
            _pick_message_for_switch_player(player)


class SwitchMessage(Page):
    @staticmethod
    def is_displayed(player: Player):
        return player.switch_condition == 1


class Decision(Page):
    form_model = 'player'
    form_fields = ['roll_choice']

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            switch_condition=player.switch_condition
        )


class Recall(Page):
    form_model = 'player'
    form_fields = ['password_recall']

    @staticmethod
    def is_displayed(player: Player):
        return player.load_condition == 1

class RecallCheck(Page):
    form_model = 'player'
    form_fields = ['memory_strategy']

    @staticmethod
    def is_displayed(player: Player):
        return player.load_condition == 1
    
class Beliefs(Page):
    form_model = 'player'
    form_fields = ['belief_roll_cat']


class ThankYou(Page):
    form_model = 'player'
    form_fields = ['outcome_email']


page_sequence = [
    MyPage, Instructions, Comprehension, Message,
    LoadTask, Demographics,
    SwitchMessage, Decision,
    Recall, RecallCheck, Beliefs, ThankYou
]