from otree.api import *
import time
import random 
import itertools 

import itertools
import numpy as np
from scipy.stats import truncnorm


doc = """
Training app for the NK leader experiment:)
"""

# TKTK: Replaced for training with values
# ── Cover story atoms ────────────────────────────────────────────────────────
# Change robots or authority options here; all scenario text updates automatically.

# ROBOTS = ["Robot K3M", "Robot N7P", "Robot V2X", "Robot Z5B"] # TK: Duplicated with button_labels

# AUTHORITY_PHRASES = {
#     "Research Team to decide":      "the Research Team members to decide by majority voting",
#     "Mission Commander to decide":  "the Mission Commander to make the final decision",
# }


# ── Distribution parameters (mirror main experiment; easy to update) ──────────
# Design decision (conversation): use truncated normals centred on four tiers
# so training estimates feel like real experiment values.
ESTIMATE_MEANS   = [30, 45, 60, 75]   # one mean per robot quality tier
ESTIMATE_SD      = 5                  # standard deviation
ESTIMATE_RANGE   = 10                 # half-width of truncation: draws stay within mean ± RANGE
ESTIMATE_MIN_GAP = 5                  # minimum integer gap between any two displayed estimates
# Design decision: fixed seed → identical scenario list every session, so
# different participants in the same session see consistent practice material.
ESTIMATE_SEED    = 42

# ── Labels and mappings ───────────────────────────────────────────────────────
# Design decision (conversation): plain "Robot N" labels in training to avoid
# anchoring participants on any named robot from the main experiment.
ROBOTS     = ["Robot 1", "Robot 2", "Robot 3", "Robot 4"]
N_ROBOTS   = len(ROBOTS)

AUTHORITY_PHRASES = {
    "researcher": "the Research Team to make the final decision",
    "commander":  "the Mission Commander to make the final decision",
}
AUTHORITY_DISPLAY = {
    "researcher": "Research Team",
    "commander":  "Mission Commander",
}
# Must match the order of C.LEADER_BUTTON_LABELS in the HTML template.
AUTHORITY_IDX = {
    "researcher": 0,
    "commander":  1,
}


# ── Estimate helpers ──────────────────────────────────────────────────────────

def _sample_sorted_estimates(rng):
    """
    Draw one integer estimate per robot from a truncated normal centred on
    each tier in ESTIMATE_MEANS. Returns a list sorted ascending.
    Enforces ESTIMATE_MIN_GAP between adjacent values (upward nudge only,
    so values stay realistic relative to their tier).
    """
    a = -ESTIMATE_RANGE / ESTIMATE_SD
    b =  ESTIMATE_RANGE / ESTIMATE_SD
    raw = sorted(
        int(round(truncnorm.rvs(a, b, loc=m, scale=ESTIMATE_SD, random_state=rng)))
        for m in ESTIMATE_MEANS
    )
    for i in range(1, len(raw)):
        if raw[i] - raw[i - 1] < ESTIMATE_MIN_GAP:
            raw[i] = raw[i - 1] + ESTIMATE_MIN_GAP
    return raw  # [lowest, ..., highest]


def _assign_estimates(sorted_vals, personal_idx, estimate_aligned, rng):
    """
    Assign the four sorted estimate values to robot indices 0..N_ROBOTS-1.

    estimate_aligned=True  → personal_idx receives the highest value.
    estimate_aligned=False → personal_idx receives the second-highest value;
                             a different robot receives the highest, motivating
                             the 'vote against the top estimate' scenario.

    Design decision (conversation): misaligned personal preference is always
    second-highest (not lower) to keep the scenario plausible — participants
    should feel they have a reasonable case for preferring a robot whose
    estimate is close to, but not at, the top.
    """
    if estimate_aligned:
        personal_val = sorted_vals[-1]
        pool         = sorted_vals[:-1]
    else:
        personal_val = sorted_vals[-2]                       # second-highest
        pool         = sorted_vals[:-2] + [sorted_vals[-1]]  # two lower + highest

    other_idxs    = [i for i in range(N_ROBOTS) if i != personal_idx]
    pool_shuffled = rng.permutation(pool).tolist()

    return {personal_idx: personal_val, **dict(zip(other_idxs, pool_shuffled))}


# ── Scenario text builder ─────────────────────────────────────────────────────

def _build_scenario_text(p_idx, t_idx, m_idx, auth_phrase, estimates, estimate_aligned):
    """
    Produce an explicit three-part scenario narrative:
      1. Personal preference motivation (with estimate context)
      2. Team vote motivation (archetype-driven)
      3. Authority preference

    Cover story terms only — no 'site', 'terrain', or 'sensor' language.
    Misaligned scenarios name both the highest-estimate robot and the
    preferred robot with their exact values so participants can clearly see
    they are being asked to vote against the top estimate, motivated by the
    idea that weather conditions affect which robot truly performs best.

    Design decision (conversation): misaligned personal preference is always
    second-highest (not lower) to keep the scenario plausible — participants
    should feel they have a reasonable case for preferring a robot whose
    estimate is close to, but not at, the top.
    """
    p = ROBOTS[p_idx]
    t = ROBOTS[t_idx]
    m = ROBOTS[m_idx]

    p_val       = estimates[p_idx]
    highest_idx = max(estimates, key=estimates.get)
    highest_rob = ROBOTS[highest_idx]
    highest_val = estimates[highest_idx]

    # Part 1 — personal preference
    if estimate_aligned:
        pref = (
            f"{p} has the highest estimated sample quality among the four robots "
            f"({p_val} points). Based on previous missions, "
            f"you trust this estimate and personally prefer {p}."
        )
    else:
        # Design decision (conversation): weather conditions are the stated
        # reason estimates can mislead, consistent with the cover story framing
        # in the introduction — no new terms introduced here.
        pref = (
            f"{highest_rob} has the highest estimated sample quality ({highest_val} points), "
            f"while {p} has a lower estimate ({p_val} points). "
            f"However, given previous missions, you believe "
            f"{p} will actually collect better samples than {highest_rob}. "
            f"You personally prefer {p}."
        )

    # Part 2 — team vote (archetype)
    if p_idx == t_idx == m_idx:       # alignment
        vote = (
            f"You are confident your teammates will also prefer {p}, "
            f"so you vote for {p} as part of the Research Team."
        )
    elif p_idx == t_idx:              # signaling
        vote = (
            f"You expect your teammates to vote for {m} instead. "
            f"Even so, you vote for {p} as part of the Research Team "
            f"to signal that {p} may be the better choice this round."
        )
    else:                             # strategic_conformity
        vote = (
            f"You expect most of your teammates to vote for {m}. "
            f"To avoid splitting the team vote, you decide to vote for {m} "
            f"as part of the Research Team, even though you personally prefer {p}."
        )

    # Part 3 — authority
    auth = f"You also want {auth_phrase}."

    return f"{pref}<br><br>{vote} {auth}"


# ── Generator ─────────────────────────────────────────────────────────────────

def _generate_scenarios():
    """
    Cross-product of valid (p, t, m) archetype combos × authority × estimate
    alignment, with a fixed RNG seed for reproducibility.

    Valid archetype combos (design decision: unchanged from original):
      alignment:            p == t == m           →  4 combos
      signaling:            p == t, p != m        → 12 combos
      strategic_conformity: p != t, t == m        → 12 combos
    Total: 28 combos × 2 authority × 2 alignment = 112 scenarios.
    Participants see 3 scenarios each, so variety is ample.
    """
    rng = np.random.default_rng(ESTIMATE_SEED)
    scenarios = []

    valid_combos = [
        (p, t, m)
        for p, t, m in itertools.product(range(N_ROBOTS), repeat=3)
        if (p == t == m) or (p == t and p != m) or (p != t and t == m)
    ]

    for (p, t, m), (auth_key, auth_phrase), estimate_aligned in itertools.product(
        valid_combos,
        AUTHORITY_PHRASES.items(),
        [True, False],
    ):
        sorted_vals = _sample_sorted_estimates(rng)
        estimates   = _assign_estimates(sorted_vals, p, estimate_aligned, rng)
        text        = _build_scenario_text(p, t, m, auth_phrase, estimates, estimate_aligned)

        if   p == t == m: archetype = "alignment"
        elif p == t:      archetype = "signaling"
        else:             archetype = "strategic_conformity"

        scenarios.append({
            "text":                text,
            "personal_preference": p,         # int robot index
            "team_choice":         t,         # int robot index
            "authority_choice":    auth_key,  # "researcher" | "commander"
            "archetype":           archetype,
            "estimate_aligned":    estimate_aligned,
            "estimates":           estimates, # {robot_idx: int}
        })

    return scenarios

SCENARIOS = _generate_scenarios()


# ── Template variable supplier ────────────────────────────────────────────────

class C(BaseConstants):

    NAME_IN_URL = 'training'
    PLAYERS_PER_GROUP = None

    NUM_QU_REPEATS = 3
    NUM_ROUNDS = 50 # Ensure all possible loops required # TK function for skipping to the end


    N_STRATEGIES = 4
    BUTTON_LABELS =  ["Robot K3M", "Robot N7P", "Robot V2X", "Robot Z5B"]
   

    LEADER_BUTTON_LABELS = ["Research Team", "Mission Commander"]

    # TK: 10 cents to 38 cents for 0 to 100 score with max -15cents fine hits minimum (plus training buffer)
    # Think it would be nice to make the price populated dynamic but not today 26.09

    # TK: Could be made dynamic with the setting value from score 

    CORRECT_CLASS_PAY_ANS = {
        'A': {'class_avg': 25, 'payoff': 17},
        'B': {'class_avg': 50, 'payoff': 22},
        'C': {'class_avg': 75, 'payoff': 28},
        'D': {'class_avg': 90, 'payoff': 34}
    }
    
    # Incorrect class performance values
    INCORRECT_CLASS_PAY_ANS = {
        'E': {'class_avg': 30, 'payoff': 8},
        'F': {'class_avg': 60, 'payoff': 10},
        'G': {'class_avg': 40, 'payoff': 7.5},
        'H': {'class_avg': 80, 'payoff': 14},
        'I': {'class_avg': 35, 'payoff': 9.5},
        'J': {'class_avg': 70, 'payoff': 11.5},
        'K': {'class_avg': 45, 'payoff': 6},
        'L': {'class_avg': 85, 'payoff': 13.5}
    }
    
    # Working with your original structure - keeping it exactly as you had it
    CORRECT_SPEED_PAY_ANS = {
        'A': {'1 Round of Voting': 1, 'payoff': 0},
        'B': {'2 Rounds of Voting': 2, 'payoff': -5},
        'C': {'3 Rounds of Voting': 3, 'payoff': -10},
        'D': {'4 Rounds of Voting': 4, 'payoff': -15}  # Fixed: was 90, should be 4
    }
    
    # Incorrect speed payment values - using same structure
    INCORRECT_SPEED_PAY_ANS = {
        'E': {'1 Round of Voting': 1, 'payoff': 8},
        'F': {'2 Rounds of Voting': 2, 'payoff': 9},
        'G': {'3 Rounds of Voting': 3, 'payoff': 6.5},
        'H': {'4 Rounds of Voting': 4, 'payoff': 4},
        'I': {'1 Round of Voting': 1, 'payoff': 11.5},
        'J': {'2 Rounds of Voting': 2, 'payoff': 6},
        'K': {'3 Rounds of Voting': 3, 'payoff': 6.5},
        'L': {'4 Rounds of Voting': 4, 'payoff': 3.5}
    }

    # Generated outside but populates the C class 
    SCENARIOS = SCENARIOS
    

# New attention checks 
    ATTENTION_CHECKS = [
    {
        'question': 'Which of these does not fit in a suitcase?',
        'options': ['Apple', 'Watch', 'Rocks', 'Car'],
        'correct': 4  # Car
    },
    {
        'question': 'Which of these cannot fly?',
        'options': ['Bird', 'Airplane', 'Elephant', 'Butterfly'],
        'correct': 3  # Elephant
    },
    {
        'question': 'Which of these is not a color?',
        'options': ['Red', 'Blue', 'Tuesday', 'Green'],
        'correct': 3  # Tuesday
    },
    {
        'question': 'Which of these cannot be eaten?',
        'options': ['Apple', 'Bread', 'Hammer', 'Cheese'],
        'correct': 3  # Hammer
    },
    {
        'question': 'Which of these is largest?',
        'options': ['Ant', 'Mouse', 'Elephant', 'Cat'],
        'correct': 3  # Elephant
    },
    {
        'question': 'Which of these is not a fruit?',
        'options': ['Orange', 'Banana', 'Carrot', 'Apple'],
        'correct': 3  # Carrot
    }
]
    
    # Standardise the question text 

    QUESTIONS_TRIPLE = ["Practice Question 1: What robot do you personally think would be best for this round?",
                        "🤖 Practice Question 2: What robot would you choose as part of the Research Team for this round?",
                        "🧑‍🔬🎖️ Practice Question 3: Who should decide which type of robot to send for this round?"]


class Subsession(BaseSubsession):
    pass


# Clean this up for differnet participant variables
def creating_session(subsession):
    session = subsession.session
    # counter for keepin track of treatments
    session.bribed_sim_com = 0
    session.bribed_com_sim = 0
    session.not_bribed_sim_com = 0
    session.not_bribed_com_sim = 0

    for player in subsession.get_players():
        player.participant.active = True # records who's inactive 
        
        player.participant.vars['strategy_tally'] = {'V0': {}, 'V1': {}, 'V2': {}}
        player.participant.vars['payoffs_cache']  = []

        # Need to time them from finishing the training
        player.participant.pay_APP = False
        player.participant.voting_APP = False
        player.participant.orphaned = False
        player.participant.is_single = False

        # Sophisticated inactivty 
        player.participant.forced_submit = False # This goes back and forth in data 
        player.participant.local_forced_submits = 0
        player.participant.global_forced_submits = 0 
        player.participant.resurrections = 0
        player.participant.resurrect_me = False
        player.participant.admonish = False
        player.participant.votes = 0
        player.participant.which_inactive = -99

        # 
        player.participant.personal_answer = -1


class Group(BaseGroup):
    pass

# Less readable but more modular code from index values corresponding to html text 
class Player(BasePlayer):

    committed = models.BooleanField()

    # Right now I do not use any of these questions as I do not collect demographics 
    consented = models.BooleanField(
        label="I acknowledge that I have read and understood the information above and confirm that I wish to participate in this study.")
    age = models.IntegerField(label='How old are you?', 
                              min=18, 
                              max=100)
    gender = models.StringField(choices=[['Man', 'Man'], 
                                         ['Woman', 'Woman'],
                                         ['Other', 'Other'],
                                         ['Prefer not to say', 'Prefer not to say']],
                                         label='What is your gender?',
                                         widget=widgets.RadioSelectHorizontal())
    
    education_lvl = models.StringField(
        choices=[['Less than high school', 'Less than high school'],
                ['High school diploma or equivalent', 'High school diploma or equivalent'],
                ['University degree or more', 'University degree or more']],
        label='What is the highest level of education you have completed?',
        widget=widgets.RadioSelect)
    

    # PAY QUIZ
    # TK comment out
    # pay_quiz = models.IntegerField(
    # choices=[
    #     [1, 'Placeholder 1 - not changed'],
    #     [2, 'Placeholder 2 - not changed'], 
    #     [3, 'Placeholder 3 - not changed'],
    #     [4, 'Placeholder 4 - not changed']
    # ],
    # widget=widgets.RadioSelect)   
    pay_was_wrong = models.BooleanField(initial=False) 
    
    pay_attended = models.BooleanField() # True or False, default to None 
    pay_trained = models.BooleanField()

    # INTERFACE QUIZ
    personal_answer = models.IntegerField() # TK personal votes
    strategy_answer = models.IntegerField()
    governance_answer = models.IntegerField()
    vote_was_wrong = models.BooleanField(initial=False)
    vote_trained = models.BooleanField()
 
    # Broccoli is a vegetable = 1 
    vote_attention = models.StringField()

    # ICON QUIZ 
    icon_quiz= models.IntegerField()
    icon_quiz_attention = models.IntegerField()
    icon_quiz_correct = models.BooleanField()
    icon_was_wrong = models.BooleanField(initial=False)
    icon_trained = models.BooleanField()

    # ATTENTION CHECK SYSTEM (new fields)
    attention_check_answer = models.IntegerField() # Come to you later for indexing 
    attention_check_index = models.IntegerField(initial=0)  # Which question from C.ATTENTION_CHECKS they got
    attention_check_count = models.IntegerField(initial=0)  # How many times they've seen attention checks
    passed_attention_check = models.BooleanField(initial=False)  # Whether they passed
    
    # PRACTICE WIDE VARIABLES # 
    attn_fail = models.BooleanField() 
    training_fail = models.BooleanField(initial = False)
    failure_shown = models.BooleanField()
    vote_qus = models.IntegerField(initial=0)
    pay_qus = models.IntegerField(initial=0) 
    icon_qus = models.IntegerField(initial=0)

    payment_quiz_answer = models.IntegerField(
    choices=[
        (1, ''),
        (2, ''),
        (3, ''),
        (4, ''),
    ],
    widget=widgets.RadioSelect,
    blank=True
)

    # Optional: to store whether they got it right
    payment_quiz_correct = models.BooleanField(initial=False)

    # Shuffle test
    shuffle_test = models.IntegerField()

# I have a load of helper functions to conditionally show pages as people loop over training to get further practice 

## FUNCTIONS ## 


def consented(player):
    return player.participant.consented

def is_pay_trained(player):
    return player.participant.vars.get('pay_trained',False)

def is_vote_trained(player):
    return player.participant.vars.get('vote_trained',False)

def is_icon_trained(player):
    return player.participant.vars.get('icon_trained', False)

# PAY FUNCTIONS # 
def should_show_pay_question(player):
    return (consented(player) and not
            is_pay_trained(player) and not
            failed_attention(player) and not 
            failed_training(player))


def check_passed_pay_training(player):
    if player.payment_quiz_answer == player.participant.correct_payment_answer_value:
        player.pay_was_wrong = False
        player.participant.pay_qus += 1
        print("Correct pay questions: ", player.participant.pay_qus)
        # Check if they've hit the threshold
        if player.participant.pay_qus >= C.NUM_QU_REPEATS:
            player.participant.pay_trained = True
        else:
            player.participant.pay_trained = False
    else:
        player.pay_was_wrong = True
        player.participant.pay_trained = False


def should_show_pay_retry(player):
    # This will work for try2 and try3 and try3 will also be conditioned on attention check separately
    return (consented(player) and
            player.pay_was_wrong and not 
            failed_attention(player) and not 
            failed_training(player)) 

# VOTE FUNCTIONS # 

def should_show_vote_question(player):
    return (consented(player) and 
            is_pay_trained(player) and not 
            is_vote_trained(player) and not
            failed_attention(player) and not 
            failed_training(player))


def check_passed_vote_training(player):
    # Get the answers submitted by the player
    player_personal = player.personal_answer
    player_strategy = player.strategy_answer
    player_governance = player.governance_answer
    
    # Get the correct answers stored in participant
    correct_personal = player.participant.correct_personal_answer
    correct_strategy = player.participant.correct_strategy_answer
    correct_governance = player.participant.correct_governance_answer
    
    # Check if each answer is correct
    personal_correct = player_personal == correct_personal
    strategy_correct = player_strategy == correct_strategy
    governance_correct = player_governance == correct_governance
    
    
    # All three answers must be correct
    all_correct = personal_correct and strategy_correct and governance_correct
    

    
    if all_correct:
        player.vote_was_wrong = False
        player.participant.vote_qus += 1

        
        # Check if they've hit the threshold
        if player.participant.vote_qus >= C.NUM_QU_REPEATS:
            player.participant.vote_trained = True

        else:
            player.participant.vote_trained = False

    else:
        player.vote_was_wrong = True
        player.participant.vote_trained = False


def should_show_vote_retry(player):
    # This will work for try2 and try3 and try3 will also be conditioned on attention check separately
    return (consented(player) and
            player.vote_was_wrong and not 
            failed_attention(player) and not 
            failed_training(player)) 

# ICON FUNCTIONS #

# Helper function to convert strategy tally to description
def tally_to_description(tally_dict):
    """Convert strategy tally to text description"""
    descriptions = []
    for strategy in range(1, C.N_STRATEGIES + 1):
        votes = tally_dict[strategy]
        if strategy == 1:
            option_name = "First option"
        elif strategy == 2:
            option_name = "Second option"  
        elif strategy == 3:
            option_name = "Third option"
        elif strategy == 4:
            option_name = "Fourth option"
        else:
            option_name = f"Option {strategy}"
        
        vote_text = "vote" if votes == 1 else "votes"
        descriptions.append(f"{option_name} received {votes} {vote_text}")
    
    return ", ".join(descriptions) + "."


def should_show_icon_question(player):
        return (consented(player) and 
            is_pay_trained(player) and 
            is_vote_trained(player) and not
            is_icon_trained(player) and not
            failed_attention(player) and not 
            failed_training(player))


def check_passed_icon_training(player):
    icon_correct = player.icon_quiz == player.participant.correct_icon_answer_value
    
    if icon_correct:
        player.icon_was_wrong = False
        player.participant.icon_qus += 1
        print("Correct icon questions: ", player.participant.icon_qus)
        # Check if they've hit the threshold
        if player.participant.icon_qus >= C.NUM_QU_REPEATS:
            player.participant.icon_trained = True
        else:
            player.participant.icon_trained = False
    else:
        player.icon_was_wrong = True
        player.participant.icon_trained = False


def should_show_icon_retry(player):
        return (consented(player) and
            player.icon_was_wrong and not 
            failed_attention(player) and not 
            failed_training(player)) 

# General functions # 

def failed_attention(player):
    return player.participant.vars.get('attn_fail',False)


def failed_training(player):
    return player.participant.vars.get('training_fail',False) 


def checked_passed_all_training(player):
    return (is_pay_trained(player) and 
            is_vote_trained(player) and 
            is_icon_trained(player)) # and blahblah


def requires_vote_training(player):
    return player.participant.vars.get('vote_trained', False)


def shuffle_button_labels(player):
    """Shuffle button labels and store as participant variables"""
    # Define the button labels and their corresponding numeric values
    labels = C.BUTTON_LABELS
    values = [1, 2, 3, 4]
    
    # Zip them together, shuffle, then unzip
    paired = list(zip(labels, values))
    random.shuffle(paired)
    
    # Unzip back into separate lists
    shuffled_labels, shuffled_values = zip(*paired)
    
    # Store as participant variables
    player.participant.vars['button_labels'] = list(shuffled_labels)
    player.participant.vars['button_values'] = list(shuffled_values)
    
    # Shuffle the "who" labels (governance question)
    who_labels = C.LEADER_BUTTON_LABELS
    who_values = [0, 1]
    
    # Zip them together, shuffle, then unzip
    who_paired = list(zip(who_labels, who_values))
    random.shuffle(who_paired)
    
    # Unzip back into separate lists
    shuffled_who_labels, shuffled_who_values = zip(*who_paired)
    
    # Store as participant variables
    player.participant.vars['button_labels_who'] = list(shuffled_who_labels)
    player.participant.vars['button_values_who'] = list(shuffled_who_values)
    
# Using base template for icon training 

# Shared vars_for_template logic extracted from Icon_try1/2/3 — all three
# classes were identical here, so a single helper avoids drift across attempts.
def _icon_vars(player):
    strategy_tally = {i: 0 for i in range(1, C.N_STRATEGIES + 1)}
    votes_remaining = 4

    while votes_remaining > 0:
        strategy = random.randint(1, C.N_STRATEGIES)
        votes_to_add = random.randint(1, votes_remaining)
        strategy_tally[strategy] += votes_to_add
        votes_remaining -= votes_to_add

    option_labels = {
        1: C.BUTTON_LABELS[0],
        2: C.BUTTON_LABELS[1],
        3: C.BUTTON_LABELS[2],
        4: C.BUTTON_LABELS[3]
    }

    def tally_to_description(tally_dict):
        descriptions = []
        for strategy in range(1, C.N_STRATEGIES + 1):
            votes = tally_dict[strategy]
            option_name = option_labels[strategy]
            vote_text = "vote" if votes == 1 else "votes"
            descriptions.append(f"{option_name}: {votes} {vote_text}")
        return "<br>".join(descriptions)

    correct_description = tally_to_description(strategy_tally)
    incorrect_descriptions = []

    for _ in range(5):
        incorrect_tally = {i: 0 for i in range(1, C.N_STRATEGIES + 1)}
        votes_remaining = 4
        while votes_remaining > 0:
            strategy = random.randint(1, C.N_STRATEGIES)
            votes_to_add = random.randint(1, votes_remaining)
            incorrect_tally[strategy] += votes_to_add
            votes_remaining -= votes_to_add
        if incorrect_tally != strategy_tally:
            incorrect_desc = tally_to_description(incorrect_tally)
            if incorrect_desc not in incorrect_descriptions and incorrect_desc != correct_description:
                incorrect_descriptions.append(incorrect_desc)

    common_incorrect_patterns = [
        {1: 1, 2: 1, 3: 1, 4: 1},
        {1: 4, 2: 0, 3: 0, 4: 0},
        {1: 0, 2: 0, 3: 0, 4: 4},
        {1: 2, 2: 2, 3: 0, 4: 0},
        {1: 0, 2: 0, 3: 2, 4: 2},
        {1: 3, 2: 1, 3: 0, 4: 0},
        {1: 1, 2: 0, 3: 0, 4: 3},
    ]

    for pattern in common_incorrect_patterns:
        if len(incorrect_descriptions) >= 3:
            break
        if pattern != strategy_tally:
            incorrect_desc = tally_to_description(pattern)
            if incorrect_desc not in incorrect_descriptions and incorrect_desc != correct_description:
                incorrect_descriptions.append(incorrect_desc)

    incorrect_descriptions = incorrect_descriptions[:3]

    answer_choices = [
        (1, correct_description),
        (2, incorrect_descriptions[0]),
        (3, incorrect_descriptions[1]),
        (4, incorrect_descriptions[2])
    ]
    random.shuffle(answer_choices)

    correct_value = next(
        value for value, description in answer_choices
        if description == correct_description
    )

    # Store on participant so try2/try3 can reuse the same scenario
    player.participant.shuffled_icon_choices = answer_choices
    player.participant.correct_icon_answer_value = correct_value
    player.participant.correct_icon_description = correct_description
    player.participant.icon_strategy_tally = strategy_tally

    return dict(
        prior_vote= strategy_tally,   # was: strategy_tally
        question="Which of the following descriptions is correct?",
        shuffled_choices=answer_choices,
        qu_num=getattr(player.participant, 'icon_qus', 0) + 1,
        qu_rep=C.NUM_QU_REPEATS,
        button_labels=C.BUTTON_LABELS,
    )


# Shared retry vars — reuses the scenario stored on participant by try1
def _icon_retry_vars(player):
    icon_strategy_tally = getattr(player.participant, 'icon_strategy_tally', {1: 1, 2: 1, 3: 1, 4: 1})
    return dict(
        prior_vote= icon_strategy_tally,  # was: icon_strategy_tally
        question="Which of the following descriptions is correct?",
        shuffled_choices=getattr(player.participant, 'shuffled_icon_choices', []),
        qu_num=getattr(player.participant, 'icon_qus', 0) + 1,
        qu_rep=C.NUM_QU_REPEATS,
        button_labels=C.BUTTON_LABELS,
    )

# ── Template variable supplier ────────────────────────────────────────────────

def _voting_vars(player):
    """
    Supply all template variables for the training voting page.
    Also writes correct answers onto player.participant for use in
    before_next_page and _voting_retry_vars.

    Design decision (conversation): SCENARIOS now stores integer robot indices
    (0-3), not robot name strings, so correct answers are set directly as ints.
    No .index() lookup against C.BUTTON_LABELS is needed or safe here.

    Design decision (conversation): button_labels supplied as plain "Robot N"
    strings, overriding C.BUTTON_LABELS, to avoid name-based anchoring.
    """
    scenario_key = random.choice(range(len(C.SCENARIOS)))
    sc           = C.SCENARIOS[scenario_key]
    estimates    = sc["estimates"]

    p_idx   = sc["personal_preference"]   # int 0-3
    t_idx   = sc["team_choice"]           # int 0-3
    auth_key = sc["authority_choice"]     # "researcher" | "commander"

    # ── Store correct answers for before_next_page ─────────────────────────────
    player.participant.correct_personal_answer   = p_idx
    player.participant.correct_strategy_answer   = t_idx
    player.participant.correct_governance_answer = AUTHORITY_IDX[auth_key]

    # ── Cache strings and estimates for retry pages ────────────────────────────
    player.participant.scenario_text            = sc["text"]
    player.participant.correct_personal_string  = ROBOTS[p_idx]
    player.participant.correct_strategy_string  = ROBOTS[t_idx]
    player.participant.correct_authority_string = AUTHORITY_DISPLAY[auth_key]
    player.participant.button_estimates         = [estimates[i] for i in range(N_ROBOTS)]

    return dict(
        scenario_text             = sc["text"],
        correct_personal_display  = ROBOTS[p_idx],
        correct_team_display      = ROBOTS[t_idx],
        correct_authority_display = AUTHORITY_DISPLAY[auth_key],
        button_labels             = [f"Robot {i + 1}" for i in range(N_ROBOTS)],
        button_estimates          = [estimates[i] for i in range(N_ROBOTS)],
        qu_num                    = player.participant.vote_qus + 1,
        qu_rep                    = C.NUM_QU_REPEATS,
    )


def _voting_retry_vars(player):
    """
    Reconstruct vars for retry pages from participant storage.
    No new scenario draw — same scenario and estimates as try1.
    """
    return dict(
        scenario_text             = getattr(player.participant, "scenario_text",            ""),
        correct_personal_display  = getattr(player.participant, "correct_personal_string",  ""),
        correct_team_display      = getattr(player.participant, "correct_strategy_string",  ""),
        correct_authority_display = getattr(player.participant, "correct_authority_string", ""),
        button_labels             = [f"Robot {i + 1}" for i in range(N_ROBOTS)],
        button_estimates          = getattr(player.participant, "button_estimates", [-99] * N_ROBOTS),
        qu_num                    = player.participant.vote_qus + 1,
        qu_rep                    = C.NUM_QU_REPEATS,
    )

 

def _pay_comp_question_text(class_data, speed_data):
    """Reconstruct question text from stored class/speed data — shared by all tries."""
    voting_description = [k for k in speed_data.keys() if 'Round' in k][0]
    voting_rounds = speed_data[voting_description]
    if voting_rounds == 4:
        return (
            f"What would your total payoff be if the sample quality is "
            f"{class_data['class_avg']} with no decision (a random robot was sent)?"
        )
    return (
        f"What would your total payoff be if the sample quality is "
        f"{class_data['class_avg']} with {voting_rounds} vote(s)?"
    )


def _pay_comp_vars(player):
    """
    Generate a fresh payment question and store everything on participant
    so retry pages can reuse the same scenario without regenerating.
    """
    class_key = random.choice(list(C.CORRECT_CLASS_PAY_ANS.keys()))
    speed_key = random.choice(list(C.CORRECT_SPEED_PAY_ANS.keys()))

    class_data = C.CORRECT_CLASS_PAY_ANS[class_key]
    speed_data  = C.CORRECT_SPEED_PAY_ANS[speed_key]

    voting_description = [k for k in speed_data.keys() if 'Round' in k][0]
    voting_rounds      = speed_data[voting_description]

    correct_total_payoff = class_data['payoff'] + speed_data['payoff']

    correct_answer = {
        'class_avg':          class_data['class_avg'],
        'voting_description': voting_description,
        'voting_rounds':      voting_rounds,
        'total_payoff':       correct_total_payoff,
    }

    # Build three distinct incorrect distractors.
    # DESIGN DECISION (conversation): deduplicate on total_payoff only, not the
    # full dict, because the HTML only ever shows total_payoff to the participant.
    # Seeding seen_totals with the correct answer prevents a distractor from
    # accidentally matching the correct total and appearing as a duplicate option.
    incorrect_answers = []
    seen_totals = {correct_answer['total_payoff']}

    for _ in range(3):
        use_incorrect_class = random.choice([True, False])
        use_incorrect_speed = random.choice([True, False])
        if not use_incorrect_class and not use_incorrect_speed:
            use_incorrect_class = True

        ic_data = (
            C.INCORRECT_CLASS_PAY_ANS[random.choice(list(C.INCORRECT_CLASS_PAY_ANS.keys()))]
            if use_incorrect_class else class_data
        )
        is_data = (
            C.INCORRECT_SPEED_PAY_ANS[random.choice(list(C.INCORRECT_SPEED_PAY_ANS.keys()))]
            if use_incorrect_speed else speed_data
        )

        iv_desc   = [k for k in is_data.keys() if 'Round' in k][0]
        iv_rounds = is_data[iv_desc]
        candidate = {
            'class_avg':          ic_data['class_avg'],
            'voting_description': iv_desc,
            'voting_rounds':      iv_rounds,
            'total_payoff':       ic_data['payoff'] + is_data['payoff'],
        }
        if candidate['total_payoff'] not in seen_totals:
            incorrect_answers.append(candidate)
            seen_totals.add(candidate['total_payoff'])

    # Top up with fallback distractors if randomisation produced duplicates
    while len(incorrect_answers) < 3:
        ic_data = C.INCORRECT_CLASS_PAY_ANS[random.choice(list(C.INCORRECT_CLASS_PAY_ANS.keys()))]
        is_data = C.INCORRECT_SPEED_PAY_ANS[random.choice(list(C.INCORRECT_SPEED_PAY_ANS.keys()))]
        iv_desc = [k for k in is_data.keys() if 'Round' in k][0]
        candidate = {
            'class_avg':          ic_data['class_avg'],
            'voting_description': iv_desc,
            'voting_rounds':      is_data[iv_desc],
            'total_payoff':       ic_data['payoff'] + is_data['payoff'],
        }
        if candidate['total_payoff'] not in seen_totals:
            incorrect_answers.append(candidate)
            seen_totals.add(candidate['total_payoff'])

    answer_choices = [
        (1, correct_answer),
        (2, incorrect_answers[0]),
        (3, incorrect_answers[1]),
        (4, incorrect_answers[2]),
    ]
    random.shuffle(answer_choices)

    correct_value = next(v for v, d in answer_choices if d == correct_answer)

    # Store on participant so retry pages can reuse without regenerating
    player.participant.shuffled_payment_choices       = answer_choices
    player.participant.correct_payment_answer_value   = correct_value
    player.participant.correct_class_key              = class_key
    player.participant.correct_speed_key              = speed_key
    player.participant.correct_total_payoff           = correct_total_payoff

    return dict(
        question         = _pay_comp_question_text(class_data, speed_data),
        shuffled_choices = answer_choices,
        class_pay_dict   = C.CORRECT_CLASS_PAY_ANS,
        speed_pay_dict   = C.CORRECT_SPEED_PAY_ANS,
        qu_num           = player.participant.pay_qus + 1,
        qu_rep           = C.NUM_QU_REPEATS,
    )

def _pay_comp_retry_vars(player):
    """
    Reconstruct vars for retry pages from participant storage — no new random draw.
    Prepends 'Try again - ' to the question text to match old retry behaviour.
    """
    class_key = getattr(player.participant, 'correct_class_key', None)
    speed_key = getattr(player.participant, 'correct_speed_key', None)

    if class_key and speed_key:
        question = 'Try again - ' + _pay_comp_question_text(
            C.CORRECT_CLASS_PAY_ANS[class_key],
            C.CORRECT_SPEED_PAY_ANS[speed_key],
        )
    else:
        question = 'Try again - What would your total payoff be for the given scenario?'

    return dict(
        question         = question,
        shuffled_choices = getattr(player.participant, 'shuffled_payment_choices', []),
        class_pay_dict   = C.CORRECT_CLASS_PAY_ANS,
        speed_pay_dict   = C.CORRECT_SPEED_PAY_ANS,
        qu_num           = player.participant.pay_qus + 1,
        qu_rep           = C.NUM_QU_REPEATS,
    )

def _pay_comp_before_next(player):
    """Shared answer-checking logic for all three tries."""
    player.payment_quiz_correct = (
        player.payment_quiz_answer == player.participant.correct_payment_answer_value
    )
    check_passed_pay_training(player)
    print(f"Player answered: {player.payment_quiz_answer}")
    print(f"Correct answer:  {player.participant.correct_payment_answer_value}")
    print(f"Correct:         {player.payment_quiz_correct}")
    print(f"Expected payoff: {player.participant.correct_total_payoff}")
## PAGES ##  

# In your __init__.py or pages.py

class Consent(Page):
    form_model = 'player'
    form_fields = ['consented']

    @staticmethod
    def is_displayed(player:Player):
        if player.round_number == 1:
            return True # ask them once for consent
    
    # If you do not consent you are not active (and by default do not go to pay_APP)
    @staticmethod
    def before_next_page(player, timeout_happened):
        player.participant.consented = player.consented
        if player.participant.consented == False:
            player.participant.active = player.participant.consented
        
        # Initialise the question counters for the looping 
        # Do it on consent page so it only happens in Round 1 
        player.participant.pay_qus = 0
        player.participant.vote_qus = 0
        player.participant.icon_qus = 0

        # Tk Testing shuffle 
        shuffle_button_labels(player)


class Introduction(Page):
    def is_displayed(player:Player):
        return consented(player) and player.round_number == 1


class PayExplanation(Page):

    def is_displayed(player:Player):
        return consented(player) and player.round_number == 1 # Explain pay once
        

class Attn_Check(Page):
    form_model = 'player'
    form_fields = ['attention_check_answer']  # renamed field

    @staticmethod
    def is_displayed(player: Player):
        return (consented(player) and
                (player.pay_was_wrong or player.vote_was_wrong or player.icon_was_wrong) and
                not failed_attention(player) and
                not failed_training(player))   # added guard

    @staticmethod
    def vars_for_template(player: Player):
        check_index = player.attention_check_count % len(C.ATTENTION_CHECKS)
        player.attention_check_index = check_index
        player.attention_check_count += 1
        check = C.ATTENTION_CHECKS[check_index]
        return dict(
            question=check['question'],
            option1=check['options'][0],
            option2=check['options'][1],
            option3=check['options'][2],
            option4=check['options'][3],
            attempt_number=player.attention_check_count
        )

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        check = C.ATTENTION_CHECKS[player.attention_check_index]
        if player.attention_check_answer != check['correct']:  # renamed field
            player.participant.vars['attn_fail'] = True


class Attn_Fail(Page):

    @staticmethod
    def is_displayed(player):
        return failed_attention(player)


class Training_Fail(Page):
    @staticmethod
    def is_displayed(player):
        # Only show if they failed training AND haven't seen this page yet
        return (failed_training(player) and 
                not player.participant.vars.get('fail_shown', False))
    
    @staticmethod
    def before_next_page(player, timeout_happened):
        # Mark that they've seen the training fail page - using .vars to match the check above
        player.participant.fail_shown = True
        player.participant.committed = False
        player.participant.pay_APP = False


## TK: With values on buttons now 
# ── Page class (unchanged structure, updated vars) ────────────────────────────

class Voting_try1(Page):
    form_model  = 'player'
    form_fields = ['personal_answer', 'strategy_answer', 'governance_answer']

    @staticmethod
    def is_displayed(player: Player):
        return should_show_vote_question(player)

    @staticmethod
    def vars_for_template(player):
        return _voting_vars(player)

    @staticmethod
    def before_next_page(player, timeout_happened):
        personal_correct   = player.personal_answer   == player.participant.correct_personal_answer
        strategy_correct   = player.strategy_answer   == player.participant.correct_strategy_answer
        governance_correct = player.governance_answer == player.participant.correct_governance_answer
        player.vote_was_wrong = not (personal_correct and strategy_correct and governance_correct)
        check_passed_vote_training(player)



class Voting_try2(Page):
    form_model  = 'player'
    form_fields = ['personal_answer', 'strategy_answer', 'governance_answer']

    @staticmethod
    def is_displayed(player: Player):
        return should_show_vote_question(player) and player.vote_was_wrong

    @staticmethod
    def vars_for_template(player):
        # Reuses scenario stored by try1 — no new random draw
        return _voting_retry_vars(player)

    @staticmethod
    def before_next_page(player, timeout_happened):
        personal_correct   = player.personal_answer   == player.participant.correct_personal_answer
        strategy_correct   = player.strategy_answer   == player.participant.correct_strategy_answer
        governance_correct = player.governance_answer == player.participant.correct_governance_answer

        player.vote_was_wrong = not (personal_correct and strategy_correct and governance_correct)
        check_passed_vote_training(player)


class Voting_try3(Page):
    form_model  = 'player'
    form_fields = ['personal_answer', 'strategy_answer', 'governance_answer']

    @staticmethod
    def is_displayed(player: Player):
        return should_show_vote_retry(player)

    @staticmethod
    def vars_for_template(player):
        return _voting_retry_vars(player)

    @staticmethod
    def before_next_page(player, timeout_happened):
        personal_correct   = player.personal_answer   == player.participant.correct_personal_answer
        strategy_correct   = player.strategy_answer   == player.participant.correct_strategy_answer
        governance_correct = player.governance_answer == player.participant.correct_governance_answer

        player.vote_was_wrong = not (personal_correct and strategy_correct and governance_correct)
        check_passed_vote_training(player)

        # After the final attempt, any remaining failure becomes a permanent training fail
        if not player.participant.vote_trained:
            player.participant.vars['training_fail'] = True


class Icon_try1(Page):
    form_model = 'player'
    form_fields = ['icon_quiz']

    @staticmethod
    def is_displayed(player: Player):
        return (consented(player) and
                is_pay_trained(player) and
                is_vote_trained(player) and
                should_show_icon_question(player))

    @staticmethod
    def vars_for_template(player):
        return _icon_vars(player)

    @staticmethod
    def before_next_page(player, timeout_happened):
        player_answer = player.icon_quiz
        correct_answer = player.participant.correct_icon_answer_value

        player.icon_quiz_correct = (player_answer == correct_answer)
        player.icon_was_wrong = not player.icon_quiz_correct

        check_passed_icon_training(player)

        print(f"Player answered: {player_answer}")
        print(f"Correct answer: {correct_answer}")
        print(f"Icon quiz correct: {player.icon_quiz_correct}")


class Icon_try2(Page):
    form_model = 'player'
    form_fields = ['icon_quiz']

    @staticmethod
    def is_displayed(player):
        return should_show_icon_retry(player)

    @staticmethod
    def vars_for_template(player):
        return _icon_retry_vars(player)

    @staticmethod
    def before_next_page(player, timeout_happened):
        player_answer = player.icon_quiz
        correct_answer = player.participant.correct_icon_answer_value

        icon_correct = (player_answer == correct_answer)
        player.icon_was_wrong = not icon_correct

        check_passed_icon_training(player)


class Icon_try3(Page):
    form_model = 'player'
    form_fields = ['icon_quiz']

    @staticmethod
    def is_displayed(player):
        return should_show_icon_retry(player)

    @staticmethod
    def vars_for_template(player):
        return _icon_retry_vars(player)

    @staticmethod
    def before_next_page(player, timeout_happened):
        # Final attempt — mirrors try1/try2 correct-answer recording,
        # but additionally flags training_fail on incorrect (unlike try1/try2)
        player_answer = player.icon_quiz
        correct_answer = player.participant.correct_icon_answer_value
        icon_correct = (player_answer == correct_answer)
        player.icon_was_wrong = not icon_correct

        # Delegate to shared helper, same as try1/try2, so correct answers
        # are recorded identically regardless of which attempt succeeded
        check_passed_icon_training(player)

        # Only try3 sets training_fail — exhausted all attempts without success
        if not icon_correct:
            player.participant.vars['training_fail'] = True


class Pay_Comp_try1(Page):
    form_model  = 'player'
    form_fields = ['payment_quiz_answer']

    @staticmethod
    def is_displayed(player: Player):
        return should_show_pay_question(player)

    @staticmethod
    def vars_for_template(player):
        return _pay_comp_vars(player)

    @staticmethod
    def before_next_page(player, timeout_happened):
        _pay_comp_before_next(player)


class Pay_Comp_try2(Page):
    form_model  = 'player'
    form_fields = ['payment_quiz_answer']

    @staticmethod
    def is_displayed(player: Player):
        return should_show_pay_retry(player)

    @staticmethod
    def vars_for_template(player):
        return _pay_comp_retry_vars(player)

    @staticmethod
    def before_next_page(player, timeout_happened):
        _pay_comp_before_next(player)

class Pay_Comp_try3(Page):
    form_model  = 'player'
    form_fields = ['payment_quiz_answer']

    @staticmethod
    def is_displayed(player: Player):
        # Displayed only if the player failed try2 as well
        return should_show_pay_retry(player)

    @staticmethod
    def vars_for_template(player):
        return _pay_comp_retry_vars(player)

    @staticmethod
    def before_next_page(player, timeout_happened):
        _pay_comp_before_next(player)
        # Final attempt — failure here is permanent, matching Voting_try3 / Icon_try3 behaviour
        if not player.participant.pay_trained:
            player.participant.vars['training_fail'] = True


class Commitment(Page):
    form_model = 'player'
    form_fields = ['committed']

    # Checking pay loop 
    @staticmethod
    def is_displayed(player):
        return (consented(player) and 
                is_pay_trained(player) and 
                checked_passed_all_training(player) and
                player.participant.vars.get('committed') is None) # Only ask for committment once
    
    # @staticmethod
    # def is_displayed(player:Player):
    #     return consented(player) and requires_pay_quiz(player) and requires_vote_training(player) 
    
    @staticmethod
    def before_next_page(player, timeout_happened):

        player.participant.committed = player.committed
        
        player.participant.voting_APP = player.participant.committed # if commit go to voting app 
        player.participant.pay_APP = not player.participant.committed # if not commit go to pay app 
        player.participant.active = player.participant.committed # if you are committed we mark you as active

        player.participant.training_fail = False # If you reached commitment then you did not fail training 
        

        #

        player.participant.wait_page_arrival = time.time()
        print("Here is the wait page arrival:",player.participant.wait_page_arrival)

        # Checking training only people get the participant fee 
        print("Here is their total calculated payoff: ", player.participant.payoff_plus_participation_fee())


# Can put the Training_Fail anywhere as it will loop back to it - wrong 
# Put fails at the end so they 'know' if something has been failed 
page_sequence = [Consent,Introduction, PayExplanation, Pay_Comp_try1,Pay_Comp_try2,Attn_Check, Pay_Comp_try3, Voting_try1, Voting_try2,Attn_Check, Voting_try3,Icon_try1,Icon_try2,Attn_Check,Icon_try3, Attn_Fail, Training_Fail,Commitment]