tempCodeRunnerFile.py

import pygame
import random
import time
import sys
import re
import csv
from collections import Counter
from datetime import datetime

pygame.init()

# Screen setup
infoObject = pygame.display.Info()
WIDTH, HEIGHT = infoObject.current_w, infoObject.current_h
screen = pygame.display.set_mode((WIDTH, HEIGHT), pygame.FULLSCREEN)
pygame.display.set_caption("Free Recall Experiment")

FONT = pygame.font.SysFont("arial", 28)
SMALL_FONT = pygame.font.SysFont("arial", 22)
CLOCK = pygame.time.Clock()

# Word lists (15 words each) - provided by user
word_lists = [
    ['ember', 'topic', 'lime', 'dawn', 'fruit', 'limit', 'gate', 'crowd', 'grow', 'sheet', 'fetch', 'ties', 'cloud', 'hotel', 'track'],
    ['point', 'ally', 'mango', 'host', 'cross', 'title', 'love', 'blank', 'order', 'stage', 'drawn', 'label', 'rail', 'night', 'care'],
    ['golf', 'heart', 'wire', 'peach', 'email', 'woods', 'cart', 'smoke', 'chunk', 'sugar', 'bill', 'user', 'clock', 'chill', 'case'],
    ['stone', 'score', 'venue', 'buses', 'oats', 'plan', 'laws', 'trade', 'chat', 'great', 'chord', 'drama', 'flame', 'sport', 'knot'],
    ['text', 'admin', 'bacon', 'deal', 'tiles', 'pilot', 'post', 'brush', 'sharp', 'cars', 'pasta', 'paint', 'crime', 'trial', 'cover'],
    ['play', 'blind', 'earth', 'hour', 'lead', 'judge', 'film', 'guest', 'suite', 'props', 'plane', 'apple', 'coach', 'round', 'clear'],
    ['route', 'jury', 'taxi', 'today', 'dairy', 'drum', 'brook', 'login', 'dusk', 'bench', 'cliff', 'port', 'ship', 'code', 'soft'],
    ['scene', 'grape', 'bank', 'happy', 'calm', 'later', 'mind', 'quilt', 'chain', 'crew', 'tempo', 'road', 'gust', 'dome', 'trust'],
    ['apart', 'work', 'alone', 'storm', 'pace', 'table', 'move', 'kind', 'loud', 'drive', 'beds', 'hope', 'goal', 'awake', 'cost'],
    ['grant', 'owner', 'coast', 'alley', 'cream', 'sale', 'cozy', 'price', 'frame', 'steak', 'glass', 'badge', 'time', 'forum', 'beam'],
    ['team', 'warm', 'phone', 'noon', 'wheel', 'camp', 'chair', 'home', 'note', 'smart', 'fraud', 'ball', 'roof', 'enact', 'grass'],
    ['motor', 'reel', 'light', 'value', 'floor', 'melon', 'shelf', 'match', 'save', 'spark', 'color', 'pack', 'train', 'creek', 'axle'],
    ['grove', 'flash', 'show', 'rope', 'beat', 'vocal', 'repay', 'berry', 'clean', 'lemon', 'piece', 'door', 'jogs', 'image', 'wall'],
    ['pump', 'tire', 'legal', 'actor', 'corn', 'refs', 'alarm', 'sand', 'sleep', 'join', 'token', 'jump', 'pipes', 'data', 'early'],
    ['atlas', 'will', 'fine', 'claim', 'lobby', 'alert', 'apps', 'batch', 'tweet', 'wise', 'gear', 'tune', 'rice', 'blog', 'cast'],
    ['slow', 'porch', 'rules', 'park', 'quick', 'lane', 'eggs', 'river', 'beach', 'toast', 'ocean', 'boat', 'seat', 'slice', 'draw'],
    ['field', 'watch', 'room', 'need', 'berth', 'earn', 'frost', 'file', 'want', 'stair', 'rest', 'fast', 'trail', 'court', 'line'],
    ['tent', 'clip', 'bread', 'lock', 'brave', 'share', 'cash', 'sound', 'guide', 'cabin', 'candy', 'keys', 'club', 'lose', 'block'],
    ['short', 'quiet', 'timer', 'sign', 'game', 'dream', 'good', 'lyric', 'mirror', 'large', 'sail', 'wins', 'rock', 'hike', 'spend'],
    ['group', 'small', 'spice', 'sing', 'song', 'dock', 'soup', 'pulse', 'milk', 'sunny', 'giant', 'soul', 'scale', 'luck', 'wish'],
]

# Tasks
TASKS = ["Baseline", "Presentation Rate", "Working Memory", "Pause"]
REPEATS = 5

# Timing constants (in seconds)
INTRO_TIME = 20
INSTRUCTION_TIME = 6
PAUSE_INSTRUCTION_TIME = 4
WORD_TIME_BASELINE = 1.5
WORD_TIME_FAST = 0.75
PAUSE_AFTER_LIST = 10
RECALL_TIME = 60
SUBMIT_DISABLE_TIME = 25  # Submit disabled for first 35 seconds

# Scoring tolerance settings (tweak as needed)
SIMILARITY_THRESHOLD = 0.80  # normalized similarity >= this is considered a match
MAX_EDIT_DISTANCE = 2        # OR absolute Levenshtein distance <= this will match

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)

# Helper: format timestamp to ISO or empty string
def fmt_ts(ts):
    return datetime.fromtimestamp(ts).isoformat() if ts else ""

# Helper: draw centered text; optionally show "(Press any key to continue)" at bottom
def draw_text_center(text, font, y_offset=0, show_continue=False):
    screen.fill(WHITE)
    lines = text.split('\n')
    total_height = len(lines) * font.get_height()
    start_y = (HEIGHT - total_height) // 2 + y_offset
    for i, line in enumerate(lines):
        rendered = font.render(line, True, BLACK)
        rect = rendered.get_rect(center=(WIDTH // 2, start_y + i * font.get_height()))
        screen.blit(rendered, rect)

    if show_continue:
        cont_text = "(Press any key to continue)"
        cont_render = SMALL_FONT.render(cont_text, True, BLACK)
        cont_rect = cont_render.get_rect(center=(WIDTH // 2, HEIGHT - 40))
        screen.blit(cont_render, cont_rect)

    pygame.display.flip()

def wait_for_key():
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.KEYDOWN:
                return

def get_user_name():
    input_text = ""
    prompt = "Please enter participant's name or alias and press Enter:"
    while True:
        screen.fill(WHITE)
        # Draw prompt
        rendered_prompt = FONT.render(prompt, True, BLACK)
        screen.blit(rendered_prompt, (20, HEIGHT // 3))

        # Draw input text
        rendered_input = FONT.render(input_text, True, BLACK)
        screen.blit(rendered_input, (20, HEIGHT // 3 + 50))

        cont_render = SMALL_FONT.render("(Press Enter to continue)", True, BLACK)
        screen.blit(cont_render, (20, HEIGHT // 3 + 100))

        pygame.display.flip()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_RETURN and input_text.strip() != "":
                    return input_text.strip()
                elif event.key == pygame.K_BACKSPACE:
                    input_text = input_text[:-1]
                else:
                    if len(event.unicode) == 1 and event.unicode.isprintable():
                        input_text += event.unicode
        CLOCK.tick(30)

def settings_screen(default_fuzzy=True):
    """
    Let experimenter choose fuzzy matching (default True) or exact-match-only (False).
    Returns boolean exact_match_only.
    """
    fuzzy = default_fuzzy
    prompt_lines = [
        "Settings - scoring method",
        "",
        "Fuzzy matching (Levenshtein tolerance) allows near-miss matches.",
        "Exact-match only requires normalized tokens to match exactly.",
        "",
        f"Current mode: {'Fuzzy matching' if fuzzy else 'Exact-match only'}",
        "",
        "Press F to toggle fuzzy/exact.",
        "Press Enter to accept and continue."
    ]
    while True:
        screen.fill(WHITE)
        y = HEIGHT // 4
        for line in prompt_lines:
            rendered = FONT.render(line, True, BLACK)
            screen.blit(rendered, (40, y))
            y += FONT.get_height() + 6
        cont_render = SMALL_FONT.render("(Press Enter to continue)", True, BLACK)
        screen.blit(cont_render, (40, HEIGHT - 40))

        pygame.display.flip()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_f:
                    fuzzy = not fuzzy
                    prompt_lines[6] = f"Current mode: {'Fuzzy matching' if fuzzy else 'Exact-match only'}"
                elif event.key == pygame.K_RETURN:
                    return not fuzzy  # return exact_match_only
        CLOCK.tick(30)

def show_intro():
    intro_text = ("Thank you for participating.\n\n"
                  "In this experiment, you will hear a list of words presented one by one.\n"
                  "Your task is to remember as many words as possible.\n"
                  "After the list ends, you will be asked to write down all the words you can recall, in any order.")
    draw_text_center(intro_text, FONT, show_continue=True)
    start_time = time.time()
    while time.time() - start_time < INTRO_TIME:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.KEYDOWN:
                return
        CLOCK.tick(30)

def show_instruction(text, duration):
    draw_text_center(text, FONT, show_continue=True)
    start_time = time.time()
    while time.time() - start_time < duration:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.KEYDOWN:
                return
        CLOCK.tick(30)

def present_words(words, word_time):
    """
    Presents words (full-screen) and returns (start_ts, end_ts)
    """
    list_start = time.time()
    screen.fill(WHITE)
    pygame.display.flip()
    for word in words:
        screen.fill(WHITE)
        rendered = FONT.render(word, True, BLACK)
        rect = rendered.get_rect(center=(WIDTH // 2, HEIGHT // 2))
        screen.blit(rendered, rect)
        pygame.display.flip()
        start = time.time()
        while time.time() - start < word_time:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    sys.exit()
            CLOCK.tick(30)
    list_end = time.time()
    return list_start, list_end

def pause_blank(duration):
    """
    Show a blank white pause screen with a countdown timer.
    Do not allow skipping by key presses (only quit allowed).
    """
    start = time.time()
    while True:
        elapsed = time.time() - start
        remaining = max(0, int(duration - elapsed))

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            # ignore keypresses during pause

        screen.fill(WHITE)
        msg = "Pause - please wait"
        r = FONT.render(msg, True, BLACK)
        rect = r.get_rect(center=(WIDTH // 2, HEIGHT // 2 - 30))
        screen.blit(r, rect)

        timer_text = f"Continuing in {remaining} second(s)..."
        tr = FONT.render(timer_text, True, BLACK)
        tr_rect = tr.get_rect(center=(WIDTH // 2, HEIGHT // 2 + 30))
        screen.blit(tr, tr_rect)

        pygame.display.flip()
        CLOCK.tick(30)

        if elapsed >= duration:
            return time.time()

def wrap_text(text, font, max_width):
    words = text.split(' ')
    lines = []
    current_line = ""
    for word in words:
        test_line = current_line + word + " "
        if font.size(test_line)[0] <= max_width:
            current_line = test_line
        else:
            lines.append(current_line)
            current_line = word + " "
    lines.append(current_line)
    return lines

# Normalizes a token: lowercase, remove punctuation (keeps letters and digits)
def normalize_token(tok):
    tok = tok.lower()
    tok = re.sub(r'[^\w]', '', tok)
    return tok

# Split response into tokens by commas and whitespace, normalize, drop empties
def tokenize_response(resp_text):
    tokens = re.split(r'[,\s]+', resp_text)
    tokens = [normalize_token(t) for t in tokens if t.strip() != ""]
    return tokens

# Levenshtein distance
def levenshtein(a, b):
    a = a or ""
    b = b or ""
    la, lb = len(a), len(b)
    if la == 0:
        return lb
    if lb == 0:
        return la
    prev_row = list(range(lb + 1))
    for i, ca in enumerate(a, start=1):
        cur_row = [i] + [0] * lb
        for j, cb in enumerate(b, start=1):
            insert_cost = prev_row[j] + 1
            delete_cost = cur_row[j - 1] + 1
            replace_cost = prev_row[j - 1] + (0 if ca == cb else 1)
            cur_row[j] = min(insert_cost, delete_cost, replace_cost)
        prev_row = cur_row
    return prev_row[lb]

def score_response(expected_list, response_text, exact_match_only=False,
                   similarity_threshold=0.85,  # Increased threshold
                   max_edit_distance=1):        # Decreased distance
    """
    Scores the participant's recall response against the expected word list.
    If exact_match_only is True, only exact normalized matches are accepted.
    Otherwise, fuzzy matching uses both similarity and edit distance thresholds.
    Returns a dictionary with scoring details.
    """
    expected_norm = [normalize_token(w) for w in expected_list if w.strip() != ""]
    expected_counter = Counter(expected_norm)  # token -> remaining count

    response_tokens_raw = re.split(r'[,\s]+', response_text)
    response_tokens = [(rt.strip(), normalize_token(rt)) for rt in response_tokens_raw if rt.strip() != ""]

    matches = []     # (response_raw, response_norm, matched_expected, edit_distance, similarity)
    extras = []      # response_raw tokens that didn't match

    for response_raw, r_norm in response_tokens:
        if r_norm == "":
            continue

        if exact_match_only:
            # Only accept exact normalized token matches (consuming duplicates)
            if expected_counter.get(r_norm, 0) > 0:
                matches.append((response_raw, r_norm, r_norm, 0, 1.0))
                expected_counter[r_norm] -= 1
                if expected_counter[r_norm] <= 0:
                    del expected_counter[r_norm]
            else:
                extras.append(response_raw)
        else:
            # Fuzzy matching: find best remaining expected token by similarity and edit distance
            best_match = None
            best_sim = -1.0
            best_dist = None
            best_expected_token = None

            for expected_token, count in list(expected_counter.items()):
                if count <= 0:
                    continue
                dist = levenshtein(r_norm, expected_token)
                max_len = max(len(r_norm), len(expected_token))
                sim = 1.0 - (dist / max_len) if max_len > 0 else 1.0

                # Prioritize exact matches
                if r_norm == expected_token:
                    best_match = (response_raw, r_norm, expected_token, dist, sim)
                    best_expected_token = expected_token
                    break

                # Update best fuzzy match if it meets both thresholds
                if sim > best_sim and dist <= max_edit_distance and sim >= similarity_threshold:
                    best_sim = sim
                    best_dist = dist
                    best_expected_token = expected_token
                    best_match = (response_raw, r_norm, expected_token, dist, sim)

            if best_match:
                matches.append(best_match)
                expected_counter[best_expected_token] -= 1
                if expected_counter[best_expected_token] <= 0:
                    del expected_counter[best_expected_token]
            else:
                extras.append(response_raw)

    missed = []
    for token, count in expected_counter.items():
        missed.extend([token] * count)

    result = {
        "num_expected": len(expected_norm),
        "num_correct": len(matches),
        "matches": matches,
        "missed": missed,
        "extra": extras
    }
    return result


def get_recall_response():
    """
    Collect typed recall for RECALL_TIME seconds.
    Submit disabled for first SUBMIT_DISABLE_TIME seconds.
    Auto-submit at RECALL_TIME.
    Returns: (recall_text, recall_start_ts, recall_end_ts)
    """
    input_text = ""
    recall_start = time.time()
    start_time = recall_start
    submit_rect = pygame.Rect(WIDTH // 2 - 60, HEIGHT - 80, 120, 50)
    done = False

    while not done:
        elapsed = time.time() - start_time
        submit_enabled = elapsed >= SUBMIT_DISABLE_TIME

        # Auto-submit at RECALL_TIME
        if elapsed >= RECALL_TIME:
            done = True

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                # Enter only accepted AFTER SUBMIT_DISABLE_TIME
                if event.key == pygame.K_RETURN:
                    if submit_enabled:
                        done = True
                    # else ignore Enter
                elif event.key == pygame.K_BACKSPACE:
                    if len(input_text) > 0:
                        input_text = input_text[:-1]
                else:
                    # record printable characters during the recall interval
                    if len(event.unicode) == 1 and event.unicode.isprintable():
                        input_text += event.unicode
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if submit_enabled and submit_rect.collidepoint(event.pos):
                    done = True
                # ignore clicks on submit before time is up

        # Draw recall entry UI
        screen.fill(WHITE)
        prompt = "Now, please write down all the words you remember in any order."
        lines = prompt.split('\n')
        y = HEIGHT // 10
        for line in lines:
            rendered = FONT.render(line, True, BLACK)
            rect = rendered.get_rect(center=(WIDTH // 2, y))
            screen.blit(rendered, rect)
            y += FONT.get_height() + 5

        # Draw input box
        input_box_rect = pygame.Rect(20, HEIGHT // 3, WIDTH - 40, HEIGHT // 3)
        pygame.draw.rect(screen, BLACK, input_box_rect, 2)

        # Draw typed text with wrapping
        wrapped_text = wrap_text(input_text, FONT, input_box_rect.width - 10)
        text_y = input_box_rect.top + 5
        for line in wrapped_text:
            rendered = FONT.render(line, True, BLACK)
            screen.blit(rendered, (input_box_rect.left + 5, text_y))
            text_y += FONT.get_height() + 2

        # Draw submit button (visually disabled until time up)
        submit_color = (0, 120, 215) if submit_enabled else (180, 180, 180)
        pygame.draw.rect(screen, submit_color, submit_rect)
        submit_text = FONT.render("Submit", True, WHITE)
        submit_text_rect = submit_text.get_rect(center=submit_rect.center)
        screen.blit(submit_text, submit_text_rect)

        # Draw timer
        time_left = max(0, int(RECALL_TIME - elapsed))
        timer_text = f"Time left: {time_left} seconds"
        timer_render = SMALL_FONT.render(timer_text, True, BLACK)
        screen.blit(timer_render, (WIDTH - 320, HEIGHT - 40))

        # Draw submit enable countdown
        if not submit_enabled:
            submit_wait = max(0, int(SUBMIT_DISABLE_TIME - elapsed))
            submit_wait_text = f"Submit enabled in {submit_wait} seconds"
            sw_render = SMALL_FONT.render(submit_wait_text, True, BLACK)
            screen.blit(sw_render, (WIDTH - 400, HEIGHT - 80))

        pygame.display.flip()
        CLOCK.tick(30)

    recall_end = time.time()

    # Intermediate screen: Enter to show experimenter review
    intermediate_msg = "Recall time finished. Press Enter to show the experimenter-only review."
    draw_text_center(intermediate_msg, FONT)
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_RETURN:
                    return input_text.strip(), recall_start, recall_end
        CLOCK.tick(30)

def show_recall_results(recall_text, expected_list, exact_match_only=False):
    """
    Experimenter-only view:
      - Show participant's typed response
      - Show scoring summary computed by score_response (with duplicates & fuzzy matching)
      - Press Enter to hide response, then Enter again to continue
    Returns the computed score dict (so main() can log timestamps too)
    """
    sc = score_response(expected_list, recall_text, exact_match_only=exact_match_only,
                        similarity_threshold=SIMILARITY_THRESHOLD, max_edit_distance=MAX_EDIT_DISTANCE)

    step = 0
    while True:
        screen.fill(WHITE)
        if step == 0:
            title = "Participant's recall response (experimenter only):"
            rendered_title = FONT.render(title, True, BLACK)
            screen.blit(rendered_title, (20, 20))

            # show the typed response (wrapped)
            lines = wrap_text(recall_text, FONT, WIDTH - 40)
            y = 20 + FONT.get_height() + 10
            for line in lines:
                rendered_line = FONT.render(line, True, BLACK)
                screen.blit(rendered_line, (20, y))
                y += FONT.get_height() + 5

            prompt = "Press Enter to hide response and show scoring summary."
            rendered_prompt = SMALL_FONT.render(prompt, True, BLACK)
            screen.blit(rendered_prompt, (20, HEIGHT - 60))

        elif step == 1:
            # show scoring summary with matched pairs & distances
            title = "Scoring summary (experimenter only):"
            rendered_title = FONT.render(title, True, BLACK)
            screen.blit(rendered_title, (20, 20))

            y = 20 + FONT.get_height() + 10

            # Basic counts
            counts = [
                f"Expected words: {sc['num_expected']}",
                f"Correct (matched): {sc['num_correct']}",
                f"Extra tokens (not matched): {len(sc['extra'])}"
            ]
            for s in counts:
                rs = FONT.render(s, True, BLACK)
                screen.blit(rs, (20, y))
                y += FONT.get_height() + 6

            y += 6
            # Matched pairs
            if sc['matches']:
                header = FONT.render("Matched (response -> expected) [edit_dist, sim]:", True, BLACK)
                screen.blit(header, (20, y))
                y += FONT.get_height() + 4
                for resp_raw, resp_norm, matched, dist, sim in sc['matches']:
                    line = f"{resp_raw} -> {matched}  [{dist}, {sim}]"
                    rline = SMALL_FONT.render(line, True, BLACK)
                    screen.blit(rline, (20, y))
                    y += SMALL_FONT.get_height() + 3
                    if y > HEIGHT - 120:
                        break  # avoid overflowing; full details are saved to CSV/TXT
                y += 6

            # Missed words
            missed_line = f"Missed ({len(sc['missed'])}): {', '.join(sc['missed']) if sc['missed'] else '(none)'}"
            rmiss = FONT.render(missed_line, True, BLACK)
            screen.blit(rmiss, (20, y))
            y += FONT.get_height() + 6

            # Extras
            extra_line = f"Extra tokens ({len(sc['extra'])}): {', '.join(sc['extra']) if sc['extra'] else '(none)'}"
            rextra = FONT.render(extra_line, True, BLACK)
            screen.blit(rextra, (20, y))
            y += FONT.get_height() + 6

            prompt = "Press Enter to continue to next trial."
            rendered_prompt = SMALL_FONT.render(prompt, True, BLACK)
            screen.blit(rendered_prompt, (20, HEIGHT - 60))

        pygame.display.flip()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_RETURN:
                    step += 1
                    if step > 1:
                        return sc  # return score details for logging

def break_time():
    break_text = "Please take a short break."
    allow_continue = False
    timer_color = BLACK

    start = time.time()
    while True:
        elapsed = time.time() - start
        remaining = max(0, int(120 - elapsed))  # countdown from 120 seconds

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if allow_continue:
                    return  # continue on any key after allowed

        screen.fill(WHITE)

        # Draw break message
        rendered_break = FONT.render(break_text, True, BLACK)
        rect_break = rendered_break.get_rect(center=(WIDTH // 2, HEIGHT // 3))
        screen.blit(rendered_break, rect_break)

        # Draw timer
        timer_text = f"Time elapsed: {int(elapsed)} second(s)"
        if elapsed >= 120:
            timer_color = (255, 0, 0)  # red
        timer_render = FONT.render(timer_text, True, timer_color)
        rect_timer = timer_render.get_rect(center=(WIDTH // 2, HEIGHT // 2))
        screen.blit(timer_render, rect_timer)

        # Additional messages based on elapsed time
        if elapsed >= 120:
            msg = "It's time to continue! Please press any key when you are ready."
            rendered_msg = SMALL_FONT.render(msg, True, (255, 0, 0))
            rect_msg = rendered_msg.get_rect(center=(WIDTH // 2, HEIGHT // 2 + 50))
            screen.blit(rendered_msg, rect_msg)
            allow_continue = True
        elif elapsed >= 40:
            msg = "Press any key to continue."
            rendered_msg = SMALL_FONT.render(msg, True, BLACK)
            rect_msg = rendered_msg.get_rect(center=(WIDTH // 2, HEIGHT // 2 + 50))
            screen.blit(rendered_msg, rect_msg)
            allow_continue = True

        pygame.display.flip()
        CLOCK.tick(30)

def run_trial(task_name, word_list, exact_match_only):
    """
    Returns: (recall_text, list_start_ts, list_end_ts, recall_start_ts, recall_end_ts)
    """
    if task_name == "Baseline":
        show_instruction("Carefully read the words.\nThey will be presented at a rate of one word every 1.5 seconds.", INSTRUCTION_TIME)
        list_start, list_end = present_words(word_list, WORD_TIME_BASELINE)
        recall_text, recall_start, recall_end = get_recall_response()
        return recall_text, list_start, list_end, recall_start, recall_end

    elif task_name == "Presentation Rate":
        show_instruction("Now, the words will be presented faster, one every 0.75 seconds.\nTry to remember as many as you can.", INSTRUCTION_TIME)
        list_start, list_end = present_words(word_list, WORD_TIME_FAST)
        recall_text, recall_start, recall_end = get_recall_response()
        return recall_text, list_start, list_end, recall_start, recall_end

    elif task_name == "Working Memory":
        show_instruction("This time, after the list is presented, you will do a short task before recalling the words.", INSTRUCTION_TIME)
        list_start, list_end = present_words(word_list, WORD_TIME_BASELINE)
        # common working-memory distractor: count backward by 3s for 15 seconds
        show_instruction("Please count backward aloud from 100 by 3s for 15 seconds.", 15)
        recall_text, recall_start, recall_end = get_recall_response()
        return recall_text, list_start, list_end, recall_start, recall_end

    elif task_name == "Pause":
        show_instruction("This time, after the list ends, there will be a 10 second pause before you recall the words.", PAUSE_INSTRUCTION_TIME)
        show_instruction("Listen carefully to the words.\nThey will be presented at a rate of one word every 1.5 seconds.", INSTRUCTION_TIME)
        list_start, list_end = present_words(word_list, WORD_TIME_BASELINE)
        pause_end_ts = pause_blank(PAUSE_AFTER_LIST)  # pause_blank returns timestamp when pause ended
        # recall begins immediately after pause; get_recall_response will set its own recall_start
        recall_text, recall_start, recall_end = get_recall_response()
        # Optionally you could record pause_end_ts too, but list_end -> pause_end_ts -> recall_start are sequential
        return recall_text, list_start, list_end, recall_start, recall_end

def save_results_txt(results, participant_name):
    with open("free_recall_results.txt", "w", encoding="utf-8") as f:
        f.write(f"Participant: {participant_name}\n\n")
        for r in results:
            f.write(f"Trial {r['trial']} - Task: {r['task']}\n")
            f.write(f"Word List: {', '.join(r['word_list'])}\n")
            f.write(f"Trial start: {fmt_ts(r.get('trial_start'))}\n")
            f.write(f"List start: {fmt_ts(r.get('list_start'))}\n")
            f.write(f"List end:   {fmt_ts(r.get('list_end'))}\n")
            f.write(f"Recall start: {fmt_ts(r.get('recall_start'))}\n")
            f.write(f"Recall end:   {fmt_ts(r.get('recall_end'))}\n")
            f.write(f"Reviewer shown: {fmt_ts(r.get('reviewer_shown'))}\n")
            f.write(f"Reviewer done:  {fmt_ts(r.get('reviewer_done'))}\n")
            f.write(f"Recall Response: {r['recall']}\n")
            sc = r.get("score")
            if sc:
                f.write(f"Scoring: {sc['num_correct']} correct of {sc['num_expected']}\n")
                f.write("Matches (response -> expected) [dist, sim]:\n")
                for resp_raw, resp_norm, matched, dist, sim in sc['matches']:
                    f.write(f"  {resp_raw} -> {matched} [{dist}, {sim}]\n")
                f.write(f"Missed ({len(sc['missed'])}): {', '.join(sc['missed'])}\n")
                f.write(f"Extra ({len(sc['extra'])}): {', '.join(sc['extra'])}\n")
            f.write("\n")

def save_results_csv(results, participant_name):
    # columns: participant,trial,task,word_list,recall_text,num_expected,num_correct,matches,missed,extra,trial_start,list_start,list_end,recall_start,recall_end,reviewer_shown,reviewer_done
    with open("free_recall_results.csv", "w", newline='', encoding="utf-8") as csvfile:
        fieldnames = [
            "participant", "trial", "task", "word_list", "recall_text",
            "num_expected", "num_correct", "matches", "missed", "extra",
            "trial_start", "list_start", "list_end", "recall_start", "recall_end",
            "reviewer_shown", "reviewer_done"
        ]
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        for r in results:
            sc = r.get("score", {})
            # format matches as semicolon-separated "resp->expected[dist,sim]"
            matches_str = "; ".join([f"{resp}->{matched}[{dist},{sim}]" for resp, rnorm, matched, dist, sim in sc.get("matches", [])])
            missed_str = "; ".join(sc.get("missed", []))
            extra_str = "; ".join(sc.get("extra", []))
            writer.writerow({
                "participant": participant_name,
                "trial": r['trial'],
                "task": r['task'],
                "word_list": "|".join(r['word_list']),
                "recall_text": r['recall'],
                "num_expected": sc.get("num_expected", 0),
                "num_correct": sc.get("num_correct", 0),
                "matches": matches_str,
                "missed": missed_str,
                "extra": extra_str,
                "trial_start": fmt_ts(r.get("trial_start")),
                "list_start": fmt_ts(r.get("list_start")),
                "list_end": fmt_ts(r.get("list_end")),
                "recall_start": fmt_ts(r.get("recall_start")),
                "recall_end": fmt_ts(r.get("recall_end")),
                "reviewer_shown": fmt_ts(r.get("reviewer_shown")),
                "reviewer_done": fmt_ts(r.get("reviewer_done"))
            })

def main():
    participant_name = get_user_name()
    exact_match_only = settings_screen(default_fuzzy=True)
    show_intro()

    all_trials = []
    for _ in range(REPEATS):
        other_tasks = TASKS.copy()
        other_tasks.remove("Baseline")
        random.shuffle(other_tasks)
        trial_order = ["Baseline"] + other_tasks
        all_trials.extend(trial_order)

    if len(all_trials) > len(word_lists):
        raise RuntimeError(f"Not enough unique word lists ({len(word_lists)}) for {len(all_trials)} trials.")
    shuffled_indices = list(range(len(word_lists)))
    random.shuffle(shuffled_indices)

    results = []
    word_list_index = 0  # <--- Initialize here

    for i, task in enumerate(all_trials):
        if i > 0:
            break_time()

        trial_text = f"Trial {i + 1} of {len(all_trials)}\nTask: {task}\nParticipant: {participant_name}"
        draw_text_center(trial_text, FONT, show_continue=True)
        wait_for_key()
        trial_start_ts = time.time()

        word_list = word_lists[word_list_index % len(word_lists)]
        word_list_index += 1

        recall_text, list_start_ts, list_end_ts, recall_start_ts, recall_end_ts = run_trial(task, word_list, exact_match_only)

        reviewer_shown_ts = time.time()
        sc = show_recall_results(recall_text, word_list, exact_match_only=exact_match_only)
        reviewer_done_ts = time.time()

        results.append({
            "trial": i + 1,
            "task": task,
            "word_list": word_list,
            "recall": recall_text,
            "score": sc,
            "trial_start": trial_start_ts,
            "list_start": list_start_ts,
            "list_end": list_end_ts,
            "recall_start": recall_start_ts,
            "recall_end": recall_end_ts,
            "reviewer_shown": reviewer_shown_ts,
            "reviewer_done": reviewer_done_ts
        })

    draw_text_center("Experiment complete.\nThank you for participating!", FONT)
    pygame.time.wait(4000)

    save_results_txt(results, participant_name)
    save_results_csv(results, participant_name)

    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()