Untitled

# Author:
# Description:
#
# Template by: Teun Mathijssen, Niels Zwemmer, Julia Dawitz


from interactive_figure import interactive_figure
import matplotlib.pyplot as plt
import pandas as pd
import os
import time 
from treisman import treisman
import random

# Import anything else you need here.


def main():
    """Run the visual search experiment"""
    # Ask participant to enter participant number
    participant_number = input("Enter participant number: ")
    
    # Run the experiment (20 correct trials per condition × 12 conditions)
    data_all_trials = run_experiment()
    
    # Save the data
    save_data(data_all_trials, participant_number)


def is_correct(target_count, key_pressed):
    """
    Test whether the right key is pressed based on whether the target is present or not
    
    :param target_count: 1 if target is present, 0 if target is absent
    :param key_pressed: The key the participant pressed 
    """

    key_pressed = key_pressed.lower() # make pressed key case-insensitive

    # input checks
    if target_count != 0 and target_count != 1:
        raise ValueError("target_count should be either 0 or 1")
    
    if key_pressed != 'l' and key_pressed != 'a':
        raise ValueError("key_pressed should be 'l' or 'a'")
    
    
    if target_count == 1: # target_count == 1 so if target is present press 'l'
        correct_key = 'l'
    else: 
        correct_key = 'a' # target_count == 0 so if target is absent press 'a' 

    return key_pressed == correct_key

# testing the function
# print(is_correct(1, 'l')) # should return True
# print(is_correct(0, 'l')) # should return False
# print(is_correct(1, 'a')) # should return False
# print(is_correct(0, 'a')) # should return True
# print(is_correct(5, 'a')) # should raise ValueError for invalid target_count
# print(is_correct(1, 'c')) # should raise ValueError for invalid key_pressed

def run_trial(vs_type, target_count, setsize):
    """
    Docstring for run_trial
    
    :param vs_type: Description
    :param target_count: Description
    :param setsize: Description
    """
    
    
    # Clear previous figure
    interactive_figure.clear()

    # Hide axes for this trial 
    plt.axis('off')

    # Show the visual search display
    treisman(target_count, vs_type, setsize)
    
    #Measure reaction time

    # while loop totdat a of l is geklikt
    valid_response = False
    start_time = time.time()

    while not valid_response: 
        interactive_figure.wait_for_interaction()
        key_pressed = interactive_figure.get_last_key_press().lower()
        

        if key_pressed in ['l', 'a']:
            valid_response = True
        stop_time = time.time()

    reaction_time = stop_time - start_time

    correct = is_correct(target_count, key_pressed)

    #Print trial data
    print(key_pressed, correct, reaction_time)

    return key_pressed, correct, reaction_time
    

# run
# print(run_trial('dcol', 1, 8)) # output

def run_experiment():
    """
    Coordinate the entire visual search experiment.
    
    Participants need to get 20 correct responses per condition.
    There are 12 conditions (4 setsizes × 3 vs_types).
    Total correct trials: 240 (12 × 20)
    Total actual trials: 240+ (includes incorrect responses)
    After every 60 trials, a 30-second pause with 5-second countdown.
    
    :return: List of dictionaries containing all trial data
    """
    
    # Display instructions
    print("\n" + "="*60)
    print("VISUAL SEARCH EXPERIMENT")
    print("="*60)
    print("\nInstructions:")
    print("- A display with symbols will appear on the screen")
    print("- Your task is to determine if the TARGET symbol is present")
    print("- Press 'L' if you think the target is PRESENT")
    print("- Press 'A' if you think the target is ABSENT")
    print("- Try to be as fast and accurate as possible")
    print("\nYou will complete 20 correct trials per condition.")
    print("There are 12 conditions total (different set sizes and search types).")
    print("="*60 + "\n")
    
    input("Press ENTER when you are ready to begin...")
    
    # Set up conditions
    setsizes = [8, 24, 40, 56]
    vs_types = ['dcol', 'dsym', 'conj']
    conditions = []
    
    for vs_type in vs_types:
        for setsize in setsizes:
            conditions.append({'vs_type': vs_type, 'setsize': setsize})
    
    # Randomize the order of conditions
    random.shuffle(conditions)
        
    # Initialize data storage
    data_all_trials = []
    total_trials_counter = 0
    
    # Main experiment loop - iterate through each condition
    for condition_idx, condition in enumerate(conditions, 1):
        vs_type = condition['vs_type']
        setsize = condition['setsize']
        correct_count = 0
        trials_in_condition = 0
        
        print(f"\n--- Condition {condition_idx}/12: {vs_type} search, set size {setsize} ---")
        print(f"Complete 20 correct trials in this condition...")
        
        # Run trials until 20 correct responses in this condition
        while correct_count < 20:
            # Run a trial
            key_pressed, correct, reaction_time = run_trial(vs_type, random.choice([0, 1]), setsize)
            
            trials_in_condition += 1
            total_trials_counter += 1
            
            # Store trial data
            trial_data = {
                'condition_number': condition_idx,
                'vs_type': vs_type,
                'setsize': setsize,
                'trial_in_condition': trials_in_condition,
                'key_pressed': key_pressed,
                'correct': correct,
                'reaction_time': reaction_time,
                'total_trials': total_trials_counter
            }
            data_all_trials.append(trial_data)
            
            # Update correct count
            if correct:
                correct_count += 1
                print(f"  Trial {trials_in_condition}: CORRECT (Correct: {correct_count}/20) RT: {reaction_time:.3f}s")
            else:
                print(f"  Trial {trials_in_condition}: incorrect (Correct: {correct_count}/20) RT: {reaction_time:.3f}s")
            
            # Pause after every 60 trials
            if total_trials_counter % 60 == 0 and total_trials_counter > 0:
                print(f"\n{'='*60}")
                print(f"BREAK TIME - 30 seconds")
                print(f"{'='*60}")
                time.sleep(25)
                for remaining in range(5, 0, -1):
                    print(f"We will start again in {remaining} seconds...")
                    time.sleep(1)
                print("Resuming experiment...\n")
        
        print(f"✓ Condition {condition_idx} completed!")
    
    print("\n" + "="*60)
    print("EXPERIMENT COMPLETED!")
    print("="*60)
    print(f"Total trials: {total_trials_counter}")
    
    return data_all_trials


def save_data(data_all_trials, participant_number):
    """Save participant data to the data directory on our drive."""
    # Convert data rows to Pandas DataFrame for easy writing and reading
    df = pd.DataFrame(data_all_trials)
    print(f"\nObtained data:\n{df}\n")

    # The name under which we will store the DataFrame on our drive
    filename = os.path.realpath(f"data/{participant_number}.csv")

    # Store the DataFrame
    df.to_csv(filename, index=False)
    print(f"Saved data in:\n{filename}\n")


if __name__ == "__main__":
    # Change our working directory to this project directory
    os.chdir(os.path.dirname(os.path.realpath(__file__)))
    print(f"\nCurrent working directory:\n{os.getcwd()}\n")

    # Ensure data directory exists
    os.makedirs("data", exist_ok=True)

    main()