task_blocks.py

# Task Class definitions
# March 2021: First version: Ladan Shahshahani  - Maedbh King - Suzanne Witt,
# Revised 2023: Bassel Arafat, Jorn Diedrichsen, Incé Husain
# Revised 2024: Caroline Nettekoven

from pathlib import Path
import pandas as pd
import numpy as np
import random
from psychopy import prefs
prefs.hardware['audioLib'] = ['sounddevice']
from psychopy import visual, sound, core, event
from pyglet.window import key
import MultiTaskBattery.utils as ut
from ast import literal_eval
from copy import deepcopy
from moviepy.audio.io.AudioFileClip import AudioFileClip
import gc
import math
import json



class Task:
    """
    Task: takes in inputs from run_experiment.py and methods
    some methods are universal across all tasks. Those methods are included in the super class Task.
    There are methods like display_instructions which are the same across most of the tasks.
    There are, however, some tasks that have different instructions from the rest (like fingerSequence).
    For those tasks, a display_instructions method is defined within the corresponding class which overrides
    the universal display_instruction method.
    Each of other classes runs a unique task given input from target files and from the Task class
    (VisualSearch, SemanticPrediction, NBack, SocialPrediction, ActionObservation).
    """

    def __init__(self, info, screen, ttl_clock, const, subj_id):

        # Pointers to Screen and experimental constants
        self.screen             = screen
        self.window             = screen.window # Shortcut to window
        self.const              = const
        self.ttl_clock          = ttl_clock  # This is a reference to the clock of the run
        self.name               = info['task_name']
        self.descriptive_name   = info['descriptive_name']
        self.code               = info['task_code']
        self.task_file          = info['task_file']
        self.feedback_type      = 'none'

        # Set the instruction text height to a value defined in the constants, or use default - useful for smaller screens
        self.const.instruction_text_height = getattr(self.const, 'instruction_text_height', None) or 1

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')

    def display_instructions(self):
        """
        displays the instruction for the task
        Most tasks have the same instructions. (Tasks that have True/False responses)
        Those tasks that have different instructions will have their own routine
        """
        true_str = f"if True press {self.const.response_keys[1]}"
        false_str = f"if False press {self.const.response_keys[2]}"

        self.instruction_text = f"{self.descriptive_name} Task\n\n {true_str} \n {false_str}"

        # 3.2 display the instruction text
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        # instr.size = 0.8
        instr_visual.draw()
        self.window.flip()
    
    def run(self):
        """Loop over trials and collects data
        Data will be stored in self.trial_data

        Returns:
            info (pd.DataFrame): _description_
        """
        self.trial_data = [] # an empty list which will be appended with the responses from each trial

        for i,trial in self.trial_info.iterrows():
            t_data = trial.copy()
            # Wait for the the start of next trial
            t_data['real_start_time'],t_data['start_ttl'],t_data['start_ttl_time'] = self.ttl_clock.wait_until(self.start_time + trial.start_time )
            # Run the trial
            t_data = self.run_trial(t_data)
            # Append the trial data
            self.trial_data.append(t_data)
        self.trial_data = pd.DataFrame(self.trial_data)
        # Calculate the feedback for the run
        acc = None
        rt = None
        if self.feedback_type[:3] == 'acc':
            acc = self.trial_data['correct'].mean()
        if self.feedback_type[-2:] == 'rt':
            rt = self.trial_data['rt'].mean()
        return acc,rt
    

    def show_progress(self, seconds_left, show_last_seconds=5, height=1, width=10, x_pos=-5, y_pos=8):
        """ Displays a progress bar for the Picture Sequence task
   Args:
        seconds_left (float): 
            The number of seconds remaining in the current trial. 
            If this value is greater than `show_last_seconds`, the progress bar is not shown.

        show_last_seconds (float, optional): 
            The duration (in seconds) over which to display the progress bar at the end of a trial. 
            Default is 5 seconds. When `seconds_left` is less than this value, the progress bar appears.

        height (float, optional): 
            The vertical size of the progress bar in PsychoPy window units. Default is 1.

        width (float, optional): 
            The horizontal size of the progress bar in PsychoPy window units. Default is 10.

        x_pos (float, optional): 
            The horizontal position of the center of the progress bar in window coordinates. 
            Negative values move it leftward. Default is -5.

        y_pos (float, optional): 
            The vertical position of the center of the progress bar in window coordinates. 
            Positive values move it upward. Default is 8.
        """
        # If we are in the last five seconds of the trial, display the remaining time
        if seconds_left < show_last_seconds:
            progress = visual.Progress(
                win=self.window, 
                progress=1-(seconds_left/show_last_seconds),
                size=(width, height),
                pos=(x_pos, y_pos),
                backColor='blue',
                barColor='black',
                borderColor='black',
                lineWidth=5,
            )
            progress.draw()

    def wait_response(self, start_time, max_wait_time, show_last_seconds=None, current_stimuli=None):
        """
        waits for a response to be made and then returns the response
        Args:
            start_time (float): the time the RT-period started
            max_wait_time (float): How long to wait maximally
        Returns:
            key (str): the key that was pressed (1-4) (0 if no key was pressed)
            rt (float): the reaction time (nan if no key was pressed)
        """
        response_made = False
        key = 0
        rt = np.nan

        while (self.ttl_clock.get_time() - start_time <= max_wait_time) and not response_made:
            self.ttl_clock.update()
            if show_last_seconds is not None:
                current_stimuli.draw()
                seconds_left = max_wait_time - (self.ttl_clock.get_time() - start_time)
                self.show_progress(seconds_left,
                                show_last_seconds=show_last_seconds,
                                y_pos=6)
                self.window.flip()
            keys=event.getKeys(keyList= self.const.response_keys, timeStamped=self.ttl_clock.clock)
            if len(keys)>0:
                response_made = True
                key_char = keys[0][0]
                key = self.const.response_keys.index(key_char) + 1
                rt = keys[0][1] - start_time
        return key, rt

    def display_trial_feedback(self, give_feedback,correct_response):
        """
        display the feedback for the current trial using the color of the fixation cross

        Args:
            give_feedback (bool):
                If true, gives informative feedback - otherwise just shows fixation cross
            correct_response (bool):
                Response was correct?, False otherwise
        """
        if give_feedback:
            if correct_response:
                self.screen.check_mark('green')
            else:
                self.screen.error_cross('red')
        else:
            self.screen.fixation_cross('white')

    def save_data(self, subj_id, run_num):
        """Saves the data to the trial data file

        Args:
            subj_id (str): Subject id to determine name
            run_num (int): Number of run - inserted as first column
        """
        self.trial_data.insert(0, 'run_num', [run_num]*len(self.trial_data))
        trial_data_file = self.const.data_dir / subj_id / f"{subj_id}_task-{self.code}.tsv"
        ut.append_data_to_file(trial_data_file, self.trial_data)

    def screen_quit(self):
        """ Checks for quit or escape key presses and quits the experiment if necessary """
        keys = event.getKeys()
        for key in keys:
            if 'q' in key or 'esc' in key:
                self.window.close()
                core.quit()

    def get_audio_from_movie(self, movie_path, sample_rate=48000):
        """Seperates the audio from the movie file and returns the audio object (for better memory handling when playing movies with sound)"""

        # Gets the movie audio
        audio_clip = AudioFileClip(movie_path)
        audio_array = audio_clip.to_soundarray(fps=sample_rate)
        audio_clip.close()
        audio = sound.Sound(audio_array,sampleRate=sample_rate, stereo=True)

        return audio

class NBack(Task):

    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')
        self.stim = []
        picture_scale = self.trial_info['picture_scale'].iloc[0] if 'picture_scale' in self.trial_info.columns else 1.0
        for stim in self.trial_info['stim']:
            stim_path = self.const.stim_dir / self.name / stim
            img = visual.ImageStim(self.window, str(stim_path))
            img.size = img.size * picture_scale
            self.stim.append(img)
        self.corr_key = [self.trial_info['key_nomatch'].iloc[0], self.trial_info['key_match'].iloc[0]]

    def display_instructions(self):
        """
        displays the instruction for the task
        """
        str1 = f"Compare image to the one shown 2 previously"
        str2 = f"if match, press {self.corr_key[1]}"
        str3 = f"if no match, press {self.corr_key[0]}"
        self.instruction_text = f"{self.descriptive_name} Task\n\n {str1} \n {str2} \n {str3}"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self,trial):
        """Runs a single trial of the nback task (after it started)
        Args:
            trial (pd.Series):
                Row of trial_info, with all information about the trial

        Returns:
            trial (pd.Series):
                Row of trial_data with all response data added
        """

        # Flush any keys in buffer
        event.clearEvents()

        # display stimulus
        self.stim[trial['trial_num']].draw()
        self.window.flip()

        # collect responses 0: no response 1-4: key pressed
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['trial_dur'])
        trial['correct'] = (trial['response'] == self.corr_key[trial['trial_type']])

        # display trial feedback
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])
        return trial

class Rest(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.name          = 'rest'

    def display_instructions(self): # overriding the display instruction routine from the parent
        self.instruction_text = 'Rest: Fixate on the cross'
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        # instr.size = 0.8
        instr_visual.draw()
        self.window.flip()

    def show_stim(self):
        # show fixation cross
        self.screen.fixation_cross()

    def run_trial(self,trial):
        # get current time (self.t0)
        self.screen.fixation_cross()
        self.ttl_clock.wait_until(self.start_time + trial['trial_dur'])
        return trial

class VerbGeneration(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')
        self.trial_info['noun'] = self.trial_info['stim'].str.strip()
        self.trial_counter = 0

    def display_instructions(self): # overriding the display instruction from the parent class

        self.instruction_text = f"{self.descriptive_name} Task \n\n Silently read the words presented.  \n\n When GENERATE is shown, silently think of verbs that go with the words."
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def show_stim(self, noun):
        """ Display a word for a fixed time. """
        stim = visual.TextStim(self.window, text=noun, pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height=2)
        stim.draw()
        self.window.flip()

    def display_generate_instruction(self):
        """ Display the 'GENERATE' instruction. """
        generate_instr = visual.TextStim(self.window, text='GENERATE', pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height=2)
        generate_instr.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Run a single trial of the VerbGeneration task. """

        # Display word
        self.show_stim(trial['noun'])

        # display GENERATE instruction at the halfway point
        if self.trial_counter == len(self.trial_info) // 2:
            self.display_generate_instruction()

        self.trial_counter += 1

        # wait for trial duration
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + trial['trial_dur'])

        # display trial feedback
        self.display_trial_feedback(give_feedback= trial['display_trial_feedback'], correct_response = None)
        return trial

class TongueMovement(Task):
    """
    Tongue movement following Buckner et al., 2022! No particular feedback.
    """
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)

    def display_instructions(self):
        self.instruction_text = f"{self.descriptive_name} Task \n\n Move your tongue left to right touching your upper premolar teeth"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Run a single trial of the tonguemovement task. """

        # draw fixation cross without flipping
        self.screen.fixation_cross(flip=False)

        # Check the trial_type and display the corresponding stimulus
        if trial['trial_type'] == 'right':
            # If trial_type is 'right', show the black circle around the fixation cross
            circle_visual = visual.Circle(self.window, radius=3, edges= 100, lineWidth = 20, fillColor=None, lineColor='black')
            circle_visual.draw()

        self.window.flip()

        # wait for trial duration
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + trial['trial_dur'])

        # display trial feedback
        self.display_trial_feedback(give_feedback= trial['display_trial_feedback'], correct_response = None)

        return trial

class AuditoryNarrative(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)

    def display_instructions(self):
        self.instruction_text = f'{self.descriptive_name} Task\n\nListen to the narrative attentively.'
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Run a single trial of the AuditoryNarrative task. """
        self.screen.fixation_cross()

        # Load and play audio stimulus for the current trial
        audio_path = self.const.stim_dir / self.name / trial['stim']
        audio_stim = sound.Sound(str(audio_path))
        audio_stim.play()

        trial_dur = audio_stim.getDuration()

        # wait for trial duration
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + trial['trial_dur'])

        # display trial feedback
        self.display_trial_feedback(give_feedback= trial['display_trial_feedback'], correct_response = None)

        return trial


class SpatialNavigation(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)

    def init_task(self):
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')
        self.corr_key = [self.trial_info['key_false'].iloc[0],self.trial_info['key_true'].iloc[0]]


    def display_instructions(self):
        start_location = self.trial_info.iloc[0]['location_1']
        end_location = self.trial_info.iloc[0]['location_2']

        self.instruction_text = (f"{self.descriptive_name} Task \n\n"
                                    f"Imagine walking around your childhood home\n"
                                    f"Start in the {start_location} – end in the {end_location}\n"
                                    f"Focus on the fixation cross")
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1],  wrapWidth=20)
        instr_visual.draw()
        self.window.flip()

    def run_trial(self,trial):
        self.screen.fixation_cross()

        # wait for trial duration
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + trial['trial_dur'])

        # display trial feedback
        self.display_trial_feedback(give_feedback= trial['display_trial_feedback'], correct_response = None)

        return trial

class TheoryOfMind(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')
        self.corr_key = [self.trial_info['key_false'].iloc[0],self.trial_info['key_true'].iloc[0]]


    def display_instructions(self):
        """
        displays the instruction for the task
        """
        task_name = visual.TextStim(self.window, text=f'{self.descriptive_name.capitalize()}', height=self.const.instruction_text_height, color=[-1, -1, -1], bold=True, pos=(0, 3))
        task_name.draw()
        str1 = f"You will read a story and decide if the answer to the question is True or False."
        str2 = f"if true, press {self.corr_key[1]}"
        str3 = f"if false, press {self.corr_key[0]}"
        self.instruction_text = f"\n\n {str1} \n\n {str2} \n {str3}"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Runs a single trial of the Theory of Mind task """

        event.clearEvents()

        height = trial.get('text_height', 1.25)
        wrapWidth=25

        # Display story
        story_clean = ' '.join(trial['story'].split('\n'))
        story_formatted = '.\n'.join(story_clean.split('. '))
        story_stim = visual.TextStim(self.window, text=story_formatted, alignHoriz='center', wrapWidth=wrapWidth, pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height=height)
        story_stim.draw()
        self.window.flip()

        # wait until story duration
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + trial['story_dur'])

        # Flush any keys in buffer
        event.clearEvents()

        # Display question
        question_stim = visual.TextStim(self.window, text=trial['question'], pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height=height, wrapWidth=25)
        question_stim.draw()
        self.window.flip()

        # collect responses 0: no response 1-4: key pressed
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['question_dur'])
        trial['correct'] = (trial['response'] == self.corr_key[trial['trial_type']])

        # display trial feedback
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])

        return trial

class DegradedPassage(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)

    def display_instructions(self):
        self.instruction_text = f'{self.descriptive_name} Task \n\nListen to the audio attentively.'
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Run a single trial of the task. """

        self.screen.fixation_cross()

        # Load and play audio stimulus for the current trial
        audio_path = self.const.stim_dir / self.name / trial['stim']
        audio_stim = sound.Sound(str(audio_path))
        audio_stim.play()

        # wait for trial duration
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + trial['trial_dur'])

        # display trial feedback
        self.display_trial_feedback(give_feedback= trial['display_trial_feedback'], correct_response = None)

        return trial

class IntactPassage(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)

    def display_instructions(self):
        self.instruction_text = f'{self.descriptive_name} Task \n\nListen to the audio attentively.'
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Run a single trial of the task. """

        self.screen.fixation_cross()

        # Load and play audio stimulus for the current trial
        audio_path = self.const.stim_dir / self.name / trial['stim']
        audio_stim = sound.Sound(str(audio_path))
        audio_stim.play()

        # wait for trial duration
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + trial['trial_dur'])

        # display trial feedback
        self.display_trial_feedback(give_feedback= trial['display_trial_feedback'], correct_response = None)

        return trial

class ActionObservation(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)

    def display_instructions(self): # overriding the display instruction from the parent class
        self.instruction_text = f"{self.descriptive_name} Task \n\n Keep your head still while watching the two clips. \n\n Try and remember the knot shown."
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Runs a single trial of the ActionObservation task """
        # Assuming that 'stim' column in trial contains the file name of the video clip
        movie_file_name = trial['stim']

        # Construct the movie file path
        movie_path = Path(self.const.stim_dir) / self.name / 'clips' / movie_file_name

        # Convert Path object to string for compatibility
        movie_path_str = str(movie_path)

        # Create a MovieStim3 object
        movie_clip = visual.MovieStim(self.window, movie_path_str, loop=False)

        while movie_clip.isFinished == False:
            movie_clip.play()
            movie_clip.draw()
            self.window.flip()
            self.ttl_clock.update()

        self.screen.fixation_cross()

        # Display trial feedback
        self.display_trial_feedback(give_feedback= trial['display_trial_feedback'], correct_response = None)

        # Flush memory: This is necessary for the script to be able to run more than 1 run. Presenting movies is very memory hungry, so do not remove!
        movie_clip.unload()
        gc.collect() # Collect garbarge

        return trial

class DemandGrid(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.square_size = 1.5
        self.feedback_type = 'acc+rt'

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')
        self.corr_key = [self.trial_info['key_left'].iloc[0],self.trial_info['key_right'].iloc[0]]

    def display_instructions(self):
        """
        displays the instruction for the task
        """
        str1 = f"You will watch the sequence of boxes that light up and then choose the correct pattern"
        str2 = f"if left, press {self.corr_key[0]}"
        str3 = f"if right, press {self.corr_key[1]}"
        self.instruction_text = f"{self.descriptive_name} Task\n\n {str1} \n {str2} \n {str3}"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def create_grid(self, sequence=None, position='center',grid_size=(3,4)):
        """Creates the grid of squares for the DemandGrid task, lighting up specific squares blue if a sequence is given,
        and positions the grid left, right, or center."""
        # Calculate offsets based on the desired position
        if position == 'left':
            offset_x = -5
        elif position == 'right':
            offset_x = 5
        else:  # center
            offset_x = 0

        # Center the grid vertically
        offset_y = 0

        grid = []
        # Create and draw the grid
        for i in range(grid_size[0]):
            row = []
            for j in range(grid_size[1]):
                # Calculate position with the offsets
                square_x = (j - grid_size[0] / 2 + 0.5) * self.square_size + offset_x
                square_y = (grid_size[1] / 2 - i - 0.5) * self.square_size + offset_y

                # Determine the fill color based on the sequence
                fill_color = 'blue' if sequence and (i, j) in sequence else 'white'

                rect = visual.Rect(self.window, width=self.square_size, height=self.square_size,
                                   pos=(square_x, square_y), lineWidth=3,
                                   lineColor='black', fillColor=fill_color)
                rect.draw()
                row.append(rect)
            grid.append(row)

        return grid

    def run_trial(self, trial):
        """Runs a single trial of the DemandGrid task with two boxes lighting up at a time"""
        # Draw the entire grid in its initial state
        if 'grid_size' in trial:
            grid_size = literal_eval(trial['grid_size'])
        else:
            grid_size = (3,4)

        # Make the code adaptable to old DemandGrid implementation
        if 'num_steps' in trial:
            num_steps = trial['num_steps']
        else:
            num_steps = 3

        step_dur = trial['sequence_dur']/num_steps
        self.grid = self.create_grid(grid_size=grid_size)
        self.window.flip()

        # Display the sequence in steps
        if 'original_sequence' in trial:
            original_sequence = literal_eval(trial['original_sequence'])
        else:
            original_sequence = literal_eval(trial['grid_sequence'])

        if 'num_steps' in trial: # new implementation
            for i in range(num_steps):
                step_sequence_name = f'original_step_{i+1}'
                step_sequence = literal_eval(trial[step_sequence_name])

                for tuple in step_sequence:
                    x, y = tuple
                    self.grid[x][y].fillColor = 'blue'

                for row in self.grid:
                    for rect in row:
                        rect.draw()
                self.window.flip()
                self.ttl_clock.wait_until(self.ttl_clock.get_time() + step_dur)

                for tuple in step_sequence:
                    x, y = tuple
                    self.grid[x][y].fillColor = 'white'

        else: # old implementation
            for i in range(0, len(original_sequence), 2):  # Iterate in steps of 2
                if i + 1 < len(original_sequence):
                    pair = [original_sequence[i], original_sequence[i + 1]]
                else:
                    pair = [original_sequence[i]]  # Handle odd-length sequences

                # Highlight positions in the current pair
                for x, y in pair:
                    self.grid[x][y].fillColor = 'blue'

                # Draw and update the window
                for row in self.grid:
                    for rect in row:
                        rect.draw()
                self.window.flip()
                self.ttl_clock.wait_until(self.ttl_clock.get_time() + step_dur)

                # Reset colors after the pair
                for x, y in pair:
                    self.grid[x][y].fillColor = 'white'

        # Flush any keys in buffer
        event.clearEvents()

        # Determine which side the correct sequence will be displayed
        correct_side = trial['correct_side']


        # # Display the original and modified sequence on the left or right side
        modified_sequence = literal_eval(trial['modified_sequence'])

        original_grid = self.create_grid(sequence=original_sequence, position=correct_side, grid_size=grid_size)
        modified_grid = self.create_grid(sequence=modified_sequence, position='left' if correct_side == 'right' else 'right', grid_size=grid_size)
        self.window.flip()

        # collect responses 0: no response 1-4: key pressed
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['question_dur'])
        trial['correct'] = (trial['response'] == self.corr_key[trial['trial_type']])

        # Provide feedback if necessary
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])

        return trial

class SentenceReading(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)

    def display_instructions(self):
        self.instruction_text = f'{self.descriptive_name} Task \n\n Read each English word and press a button when the image of a hand pressing a button is displayed'
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Run a single trial of the sentence reading task. """

        # get sentence and split into words by space
        sentence = trial['stim']
        words = sentence.split()

        #show words seqeuntially each for 450ms
        for word in words:
            word_stim = visual.TextStim(self.window, text=word, pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height=2)
            word_stim.draw()
            self.window.flip()
            self.ttl_clock.wait_until(self.ttl_clock.get_time() + 0.45)

        event.clearEvents()

        # show press button image
        button_stim = visual.ImageStim(self.window, image=str(self.const.stim_dir / self.name / 'hand_press_transparent.png'))
        button_stim.draw()
        self.window.flip()
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), 0.4)

        # show blank_transparent image
        blank_stim = visual.ImageStim(self.window, image=str(self.const.stim_dir / self.name / 'blank_transparent.png'))
        blank_stim.draw()
        self.window.flip()

        # flush any keys in buffer
        event.clearEvents()

        # not displaying crosshair to replicate the localizer

        return trial

class NonwordReading(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)

    def display_instructions(self):
        self.instruction_text = f'{self.descriptive_name} Task \n\n Read each nonword word and press a button when the image of a hand pressing a button is displayed'
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Run a single trial of the nonword reading task. """

        # get sentence and split into words by space
        sentence = trial['stim']
        words = sentence.split()

        #show words seqeuntially each for 450ms
        for word in words:
            word_stim = visual.TextStim(self.window, text=word, pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height=2)
            word_stim.draw()
            self.window.flip()
            self.ttl_clock.wait_until(self.ttl_clock.get_time() + 0.45)

        event.clearEvents()

        # show press button image
        button_stim = visual.ImageStim(self.window, image=str(self.const.stim_dir / self.name / 'hand_press_transparent.png'))
        button_stim.draw()
        self.window.flip()
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), 0.4)

        # show blank_transparent image
        blank_stim = visual.ImageStim(self.window, image=str(self.const.stim_dir / self.name / 'blank_transparent.png'))
        blank_stim.draw()
        self.window.flip()

        # clear any keys in buffer
        event.clearEvents()

        # not displaying crosshair to replicate the localizer

        return trial

class OddBall(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0],self.trial_info['key_two'].iloc[0]]

    def display_instructions(self):
        """
        displays the instruction for the task
        """
        str1 = f"Press {self.corr_key[0]} when you see a red K"
        self.instruction_text = f"{self.descriptive_name} Task\n\n {str1} \n"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Run a single trial of the oddball task. """

        current_trial =  trial['stim']

        # show stem
        if current_trial == 'red_K':
            word_stim = visual.TextStim(self.window, text='K', pos=(0.0, 0.0), color='red', units='deg', height=1.5)
        elif current_trial == 'black_K':
            word_stim = visual.TextStim(self.window, text='K', pos=(0.0, 0.0), color='black', units='deg', height=1.5)
        elif current_trial == 'red_O':
            word_stim = visual.TextStim(self.window, text='O', pos=(0.0, 0.0), color='red', units='deg', height=1.5)
        elif current_trial == 'black_O':
            word_stim = visual.TextStim(self.window, text='O', pos=(0.0, 0.0), color='black', units='deg', height=1.5)

        word_stim.draw()
        self.window.flip()
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + trial['trial_dur'])

        self.display_trial_feedback(trial['display_trial_feedback'], None)

        # Flush any keys in buffer
        event.clearEvents()

        # collect responses 0: no response 1-4: key pressed
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['iti_dur'])

        # method 2 for accuracy for oddball
        if trial['trial_type'] == 1:
            trial['correct'] = True if trial['response']!= 0 else False
        elif trial['trial_type'] == 0:
            trial['correct'] = np.nan


        return trial

class FingerSequence(Task):
    """
    Finger sequence task
    """
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0],self.trial_info['key_two'].iloc[0],self.trial_info['key_three'].iloc[0],self.trial_info['key_four'].iloc[0]]


    def display_instructions(self):
        self.instruction_text = f"{self.descriptive_name} Task \n\n Using your four fingers, press the keys in the order shown on the screen\n Use all four fingers for this task"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()


    def run_trial(self, trial):
        """ Run a single trial of the finger sequence task. """
        #clear buffer
        event.clearEvents()

        # Display the sequence
        sequence = trial['stim'].split()

        # Calculate the start position for the sequence and determine the spacing between numbers
        num_items = len(sequence)
        spacing = 2.0
        start_x = -(num_items - 1) * spacing / 2

        # Show the numbers in the sequence next to each other ( using the spacing and start_x calculated above)
        for i, number in enumerate(sequence):
            pos = (start_x + i * spacing, 0.0)  # Horizontal position is adjusted based on index
            stim = visual.TextStim(self.window, text=number, pos=pos, color='black', units='deg', height=1.5)
            stim.draw()

        self.window.flip()


        sequence_start_time = self.ttl_clock.get_time() # Needed for knowing when to stop looking for key presses
        digit_start_time = sequence_start_time # Updated with each key press for calculating RT

        rt_list = np.full(num_items,np.nan)
        correct_list = np.zeros((num_items,)) # List of booleans indicating whether each press was correct needed for overall trial accuracy
        num_presses =0
        # Initialize the color for each digit in the sequence as black
        digit_colors = ['black'] * num_items
        while self.ttl_clock.get_time() - sequence_start_time < trial['trial_dur'] and num_presses < num_items:
            self.ttl_clock.update()

            keys = event.getKeys(keyList=self.const.response_keys, timeStamped=self.ttl_clock.clock)
            if keys:
                key_char, key_press_time = keys[0]
                key = self.const.response_keys.index(key_char) + 1
                rt = key_press_time - digit_start_time
                rt_list[num_presses]=rt
                digit_start_time = key_press_time

                # Check if key pressed is correct
                correct_list[num_presses] = key == int(sequence[num_presses])

                # Update color based on correctness
                digit_colors[num_presses] = 'green' if correct_list[num_presses] else 'red'

                num_presses += 1
            # Draw all digits with their adjusted colors
            for i, (number, color) in enumerate(zip(sequence, digit_colors)):
                pos = (start_x + i * spacing, 0.0)
                stim = visual.TextStim(self.window, text=number, pos=pos, color=color, units='deg', height=1.5)
                stim.draw()

            self.window.flip()

        else:
            # If the sequence is completed, wait until the end of the trial
            self.ttl_clock.wait_until(sequence_start_time + trial['trial_dur'])

        # if any press is wrong trial['correct'] needs to be false, this is for post trial feedback
        trial['correct'] = correct_list.sum()/num_items

        if np.all(np.isnan(rt_list)):
            # calculate mean rt across presses
            trial['rt'] = np.nan

        else:
            trial['rt'] = np.nanmean(rt_list)

        # display trial feedback (for whole trial)
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct']== 1)

        return trial

class FlexionExtension(Task):
    """
    Flexion extension of toes! No particular feedback.
    """
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'None'

    def init_task(self):
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')

    def display_instructions(self):
        self.instruction_text = f"{self.descriptive_name} Task \n\n Flex and extend your right and left toes"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        real_start_time, start_ttl, start_ttl_time = self.ttl_clock.wait_until(trial['start_time'])
        # Handle the display of each action
        self.stim_pair = trial['stim'].split()
        n_rep = int(trial['trial_dur'] / (trial['stim_dur'] * 2))
        stim_actions = np.tile(self.stim_pair, n_rep)

        for action in stim_actions:
            # Display action
            stim = visual.TextStim(self.window, text=action, pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height=1.5)
            stim.draw()
            self.window.flip()

            # Wait for the duration of the stimulus
            start_time = self.ttl_clock.get_time()
            self.ttl_clock.wait_until(start_time + trial['stim_dur'])

        trial['real_start_time'] = real_start_time
        trial['start_ttl'] = start_ttl
        trial['start_ttl_time'] = start_ttl_time

        # No response is expected in this task, so return trial as is
        return trial

    #semantic prediction runs (slight bug for response feedback: last word is not synced with last word in task_file..)

class SemanticPrediction(Task):

    """
    Read a sentence and decide if the last word of the sentence makes sense. Click "3" if the last word makes sense; click "4" if not. Be as accurate and fast as possible.
    """

    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')
        self.corr_key = [self.trial_info['key_false'].iloc[0], self.trial_info['key_true'].iloc[0]]

    def display_instructions(self):
        """
        displays the instruction for the task
        """
        str1 = f"You will read a sentence and decide if the last word makes sense."
        str2 = f"If it makes sense, press {self.corr_key[1]}"
        str3 = f"if it doesn't make sense, press {self.corr_key[0]}"
        self.instruction_text = f"{self.descriptive_name} Task\n\n {str1} \n {str2} \n {str3}"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Runs a single trial of the semantic prediction task """

        height_word = 2

        event.clearEvents()

        # get sentence and split into words by space
        sentence = trial['sentence']
        words = sentence.split('|')

        #show words seqeuntially each for 800ms
        for word in words:
            word_stim = visual.TextStim(self.window, text=word, pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height=height_word)
            word_stim.draw()
            self.window.flip()
            self.ttl_clock.wait_until(self.ttl_clock.get_time() + 0.8)

        event.clearEvents()

        # Fixation cross
        self.screen.fixation_cross()
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + 0.5)
        event.clearEvents()
        
        # Display last word
        last_word_stim = visual.TextStim(self.window, text=trial['last_word'], pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height= height_word, wrapWidth=30)
        last_word_stim.draw()
        self.window.flip()

        event.clearEvents()

        # collect responses 0: no response 1-4: key pressed
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['sentence_dur'])
        trial['correct'] = (trial['response'] == self.corr_key[trial['trial_type']])

        # display trial feedback
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])

        return trial

    
class VisualSearch(Task):

    """
    Look at a screen filled with shapes and identify whether an "L" is present. Click "3" if the "L" is present; click "4" if not. Be as accurate and fast as possible.
    """

    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc'

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')

        #counter initialized in order to read whether trial_type is true or not for each trial; used in generate_trial_stimuli
        self.trial_counter = 0

        #width and height determined by trial and error: x extremity of window is 14; y extremity is 10
        screen_width = 28
        screen_height = 20

        #define number of rows + col to split screen into 24 apertures
        num_rows = 4
        num_cols = 6

        #Define aperture sizes + positions
        aperture_width = screen_width / num_cols
        aperture_height = screen_height / num_rows

        positions_x = np.linspace(-screen_width / 2 + aperture_width / 2, screen_width / 2 - aperture_width / 2, num_cols)
        positions_y = np.linspace(-screen_height / 2 + aperture_height / 2, screen_height / 2 - aperture_height / 2, num_rows)

        aperture_positions = []

        for y in positions_y:
            for x in positions_x:
                aperture_positions.append((x, y))

        self.apertures = []
        for pos in aperture_positions:
            apertures = visual.Aperture(self.window, size=40, shape = 'rectangle', pos=pos, units='norm')
            self.apertures.append(apertures)

    def generate_trial_stimuli(self, num_stimuli):

        stim_images = ['90.png','180.png','270.png','360.png']

        self.stim = []

        is_target = True #determines whether target will appear on screen

        if self.trial_info['trial_type'][self.trial_counter] == 0:
            stim_images.remove('90.png') #sets a display with no target
            is_target = False

        self.trial_counter+=1

        randomly_select_apertures = random.sample(self.apertures, num_stimuli)

        for aperture in randomly_select_apertures:
            if is_target:
                stim_current = stim_images[0] #sets a display with at least one target
                is_target = False
            else:
                stim_random_idx = random.randint(0, len(stim_images)-1)  #chooses random stimuli from stim_images list
                stim_current = stim_images[stim_random_idx]

            stim_path = self.const.stim_dir/ self.name / stim_current
            stimulus = visual.ImageStim(self.window, str(stim_path), size=(0.8,0.8))
            stimulus.setPos([aperture.pos[0], aperture.pos[1]]) #puts stimuli within apertures
            self.stim.append(stimulus) #creates list of all randomly selected stimuli

    def display_instructions(self):
        """
        displays the instruction for the task
        """

        self.corr_key = [self.trial_info['key_false'].iloc[0],self.trial_info['key_true'].iloc[0]]

        str1 = f"You will survey a series of shapes and identify whether the letter ‘L’ is present."
        str2 = f"If 'L' is present, press {self.corr_key[1]}"
        str3 = f"if 'L' is not present, press {self.corr_key[0]}"
        self.instruction_text = f"{self.descriptive_name} Task\n\n {str1} \n {str2} \n {str3}"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self,trial):
        """Runs a single trial of visual search task
        """

        # Flush any keys in buffer
        event.clearEvents()

        num_stimuli = self.trial_info.loc[trial['trial_num'], 'num_stimuli']  #indicates if trial is easy (4 stimuli) or hard (8 stimuli)

        self.generate_trial_stimuli(num_stimuli)

        self.screen.fixation_cross(flip=False)  #show fixation cross in center of screen

        # Display stimuli
        for stimulus in self.stim:
            stimulus.draw()

        self.window.flip()

        # collect responses
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['trial_dur'])
        trial['correct'] = (trial['response'] == self.corr_key[trial['trial_type']])

        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])

        return trial


class RMET(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0],self.trial_info['key_two'].iloc[0],self.trial_info['key_three'].iloc[0],self.trial_info['key_four'].iloc[0]]

    def display_instructions(self):
        task_name = visual.TextStim(self.window, text=f'{self.descriptive_name.capitalize()}', height=self.const.instruction_text_height, color=[-1, -1, -1], bold=True, pos=(0, 3))
        task_name.draw()
        self.instruction_text = ""
        if 'age' in self.task_file:
            self.instruction_text += "\n\n Choose what AGE the person is."
        elif 'emotion' in self.task_file:
            self.instruction_text += "\n\n Choose what FEELING the person has."
        else:
            self.instruction_text += "\n\n Choose which AGE or FEELING best describes the person." # General instruction for both age and emotion
        self.instruction_text += f"\n\n\n{self.trial_info['key_one'].iloc[0]}. index \t{self.trial_info['key_two'].iloc[0]}. middle\t{self.trial_info['key_three'].iloc[0]}. ring\t{self.trial_info['key_four'].iloc[0]}. pinky"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1], wrapWidth=20, pos=(0, 0))
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Runs a single trial of the Reading the Mind in the Eye (RMET) task """

        # Flush any keys in buffer
        event.clearEvents()

        # --- Eyes ---
        # Get the file name
        picture_file_name = trial['stim']
        # Construct the picture file path
        picture_path = Path(self.const.stim_dir) / self.name / 'pictures' / picture_file_name
        # Convert Pathself object to string for compatibility
        picture_path_str = str(picture_path)
        # Create an ImageStim object, explicitly centered so that the four answer
        # options (placed symmetrically above and below) straddle the image.
        picture = visual.ImageStim(self.window, str(picture_path_str), pos=(0, 0))
        # Make the picture smaller
        picture_scale = getattr(self.const, 'rmet_picture_scale', None) or 0.7
        picture.size = picture.size * picture_scale



        # --- Answers ---
        # Get the answer options and split into four strings
        answer_options = str(trial['options']).split(',')

        # Create TextStim objects for each answer option with a colored index and
        # black option text. The spatial layout and string breakup are preserved:
        # only formatting (height and color) is controlled here.
        answer_stims = []
        option_height = trial.get('option_text_height', 1.2)
        index_height = option_height * 0.85
        pos_scale = trial.get('option_position_scale', 1.0)

        # Small horizontal offset so that, visually, the options sit slightly
        # closer to the centre and the eye picture looks better centered between
        # them on wide screens.
        x_offset = -1.0

        for i, option in enumerate(answer_options):
            # 0 and 1 on top left/right; 2 and 3 on bottom left/right.
            base_x = -8 if i % 2 == 0 else 6
            base_y = 5 if i < 2 else -5
            x = base_x * pos_scale + x_offset
            y = base_y * pos_scale

            option_str = option.strip()
            index_label = f'{i+1}.'
            text_label = f'  {option_str}'

            # Draw the numeric index in a distinct color (blue-ish) to visually
            # separate it from the age/feeling text.
            index_stim = visual.TextStim(
                self.window,
                text=index_label,
                pos=(x, y),
                color=[-1, -1, 1],  # blue index on grey/white background
                height=index_height,
                alignHoriz='left',
            )

            # Draw the actual answer text in standard black, slightly to the right
            # of the index so the combined appearance matches "1.  option".
            text_stim = visual.TextStim(
                self.window,
                text=text_label,
                pos=(x + 0.8, y),
                color=[-1, -1, -1],
                height=option_height,
                alignHoriz='left',
            )

            answer_stims.extend([index_stim, text_stim])

        # Display stimuli
        picture.draw()
        for answer_stim in answer_stims:
            answer_stim.draw()
        self.window.flip()

        # collect responses 0: no response 1-4: key pressed
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['trial_dur'])
        trial['correct'] = (trial['response'] == answer_options.index(str(trial['answer']))+1)

        # display trial feedback
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])

        return trial

class PictureSequence(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0],self.trial_info['key_two'].iloc[0],self.trial_info['key_three'].iloc[0],self.trial_info['key_four'].iloc[0]]

    def display_instructions(self):
        task_name = visual.TextStim(self.window, text=f'{self.descriptive_name.capitalize()}', height=self.const.instruction_text_height, color=[-1, -1, -1], bold=True, pos=(0, 3))
        task_name.draw()
        self.instruction_text = ""
        self.instruction_text += "\n\n Find the correct chronological order of the pictures."
        self.instruction_text += f"\n\n\n{self.trial_info['key_one'].iloc[0]}. index \t{self.trial_info['key_two'].iloc[0]}. middle\t{self.trial_info['key_three'].iloc[0]}. ring\t{self.trial_info['key_four'].iloc[0]}. pinky"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1], wrapWidth=20, pos=(0, 0))
        instr_visual.draw()
        self.window.flip()

    def show_presses(self, pressed_keys, positions, last_key_press_time, width=1.4, height=7, line_width=10):
        """ Displays the presses on the screen
        Args:
            pressed_keys (list): A list of the keys that have been pressed
            positions (list): A list of the positions of the images
            width (float): The width of the images
            height (float): The height of the images
            line_width (float): The width of the border
            last_key_press_time (float): The time of the last key press
        """
        #Add a black border around the selected images
        for p, pressed_key in enumerate(pressed_keys):
            color = 'blue' if p == len(pressed_keys) - 1 and not self.ttl_clock.get_time() - last_key_press_time > 1  else 'black' #Add a green border around the last selected image if the last key press was less than 2 seconds ago
            visual.Rect(self.window, size=(width, height), pos=positions[pressed_key-1], lineColor=color, lineWidth=line_width).draw()

    def run_trial(self, trial):
        """ Runs a single trial of the Reading the Mind in the Eye (RMET) task """

        # Flush any keys in buffer
        event.clearEvents()

        # Get the file name
        picture_file_name = trial['stim']
        # Construct the picture file path
        picture_paths = [str(Path(self.const.stim_dir) / self.name / 'pictures' / f"{picture_file_name} card{n}") for n in range(1,5)]
        # Sort them in the order they should be displayed
        sequence = list(map(int, trial['sequence'].split(' ')))
        picture_paths = [picture_paths[i-1] for i in sequence]

        # Define positions for a 2x2 grid layout
        height = 10
        width = 1.4*height
        x_pos = 7.2
        y_pos = 5.2
        positions = [
            (-x_pos, y_pos),  # Top-left
            (x_pos, y_pos),   # Top-right
            (-x_pos, -y_pos), # Bottom-left
            (x_pos, -y_pos)   # Bottom-right
        ]
        # Create ImageStim objects for each picture

        pictures = [visual.ImageStim(self.window, image=path, pos=pos, size=(1.4*height, height)) for path, pos in zip(picture_paths, positions)]


        # --- Answers ---
        # Create TextStim objects for each answer option
        answer_options = ['1', '2', '3', '4']
        answer_stims = []
        for i, option in enumerate(answer_options):
            x = -x_pos-width*0.4 if i % 2 == 0 else x_pos+width*0.4
            y = y_pos+height*0.4 if i < 2 else -y_pos-height*0.4
            answer_stim = visual.TextStim(self.window, text=f'{option}', pos=(x, y), color=[-1, -1, -1], height=1.3, alignHoriz='center')
            answer_stims.append(answer_stim)

        # Calculate the start position for the sequence and determine the spacing between numbers
        num_items = len(sequence)

        # collect responses 0: no response 1-4: key pressed
        sequence_start_time = self.ttl_clock.get_time() # Needed for knowing when to stop looking for key presses
        digit_start_time = sequence_start_time # Updated with each key press for calculating RT

        rt_list = np.full(num_items,np.nan)
        correct_list = np.zeros((num_items,)) # List of booleans indicating whether each press was correct needed for overall trial accuracy
        correct_sequence = np.argsort(sequence) + 1 # correct sequence of the answers to sort the images into the right order
        num_presses =0
        pressed_keys = []
        line_width = 15

        while self.ttl_clock.get_time() - sequence_start_time < trial['trial_dur']:
            self.ttl_clock.update()

            for picture in pictures:
                picture.draw()
            for answer_stim in answer_stims:
                answer_stim.draw()

            seconds_left = trial['trial_dur'] - (self.ttl_clock.get_time() - sequence_start_time)
            self.show_progress(seconds_left,
                               show_last_seconds=5,
                               height=0.9,
                               width=width,
                               x_pos=0-width*0.5,
                               y_pos=y_pos+height*0.5+0.4)
            self.show_presses(pressed_keys, positions, digit_start_time, width, height, line_width)
            self.window.flip()

            if num_presses < num_items:
                keys = event.getKeys(keyList=self.const.response_keys, timeStamped=self.ttl_clock.clock)
                if keys:
                    key_char, key_press_time = keys[0]
                    key = self.const.response_keys.index(key_char) + 1
                    rt = key_press_time - digit_start_time
                    rt_list[num_presses]=rt
                    digit_start_time = key_press_time

                    # Check if key pressed is correct
                    correct_list[num_presses] = key == int(correct_sequence[num_presses])
                    num_presses += 1
                    pressed_keys.append(key)

        # if any press is wrong trial['correct'] needs to be false, this is for post trial feedback
        trial['correct'] = correct_list.sum()/num_items
        trial['response'] = pressed_keys

        if np.all(np.isnan(rt_list)):
            # calculate mean rt across presses
            trial['rt'] = np.nan

        else:
            trial['rt'] = np.nanmean(rt_list)

        # display trial feedback (for whole trial)
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct']==1)

        return trial


class StorySequence(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0],self.trial_info['key_two'].iloc[0],self.trial_info['key_three'].iloc[0],self.trial_info['key_four'].iloc[0]]

    def display_instructions(self):
        task_name = visual.TextStim(self.window, text=f'{self.descriptive_name.capitalize()}', height=self.const.instruction_text_height, color=[-1, -1, -1], bold=True, pos=(0, 3))
        task_name.draw()
        self.instruction_text = ""
        self.instruction_text += "\n\n Find the correct chronological order of the sentences."
        self.instruction_text += f"\n\n\n{self.trial_info['key_one'].iloc[0]}. index \t{self.trial_info['key_two'].iloc[0]}. middle\t{self.trial_info['key_three'].iloc[0]}. ring\t{self.trial_info['key_four'].iloc[0]}. pinky"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1], wrapWidth=20, pos=(0, 0))
        instr_visual.draw()
        self.window.flip()

    def show_presses(self, sentences, positions, pressed_keys, last_key_press_time, wrapWidth, text_height=1):
        """ Displays the presses on the screen
        Args:
            sentences (list): A list of the sentences to be displayed
            positions (list): A list of the positions of the sentences
            pressed_keys (list): A list of the keys that have been pressed
            last_key_press_time (float): The time of the last key press
            wrapWidth (float): The width of the text
        """


        for p, pressed_key in enumerate(pressed_keys):
            color = 'blue' if p == len(pressed_keys) - 1 and not self.ttl_clock.get_time() - last_key_press_time > 1  else 'darkgrey' # Present the stimuli in blue if the last key press was less than 2 seconds ago
            visual.TextStim(self.window, text=sentences[pressed_key-1], pos=positions[pressed_key-1], color=color, height=text_height, wrapWidth=wrapWidth).draw()



    def run_trial(self, trial):
        """ Runs a single trial of the Reading the Mind in the Eye (RMET) task """

        # Flush any keys in buffer
        event.clearEvents()

        wrapWidth = 20

        # Sort them in the order they should be displayed
        sequence = list(map(int, trial['sequence'].split(' ')))
        sentences = [trial[f"stim{i}"] for i in range(1,5)]
        sentences = [sentences[i-1] for i in sequence] # Order the sentences according to the sequence
        # Format the sentences for display
        sentences = [f'{s+1}.\t{sentence}\n\n' for s, sentence in enumerate(sentences)]
        sentences_stim = visual.TextStim(self.window, text=''.join(sentences), pos=(0, 0), color=[-1, -1, -1], height=1, wrapWidth=wrapWidth)

        # Calculate the start position for the sequence and determine the spacing between numbers
        num_items = len(sequence)

        # collect responses 0: no response 1-4: key pressed
        sequence_start_time = self.ttl_clock.get_time() # Needed for knowing when to stop looking for key presses
        digit_start_time = sequence_start_time # Updated with each key press for calculating RT

        rt_list = np.full(num_items,np.nan)
        correct_list = np.zeros((num_items,)) # List of booleans indicating whether each press was correct needed for overall trial accuracy
        num_presses =0
        pressed_keys = []
        line_width = 10
        bar_width = 10
        height = 7
        text_height = 1
        y_pos = 5

        # Arrange sentences non-overlapping from top to bottom
        positions = [(0, 5), (0, 1), (0, -2), (0, -6)]
        # Present the stimuli in black

        while self.ttl_clock.get_time() - sequence_start_time < trial['trial_dur']:
            self.ttl_clock.update()

            seconds_left = trial['trial_dur'] - (self.ttl_clock.get_time() - sequence_start_time)
            self.show_progress(seconds_left,
                               show_last_seconds=5,
                               height=1,
                               width=bar_width,
                               x_pos=0-bar_width*0.5,
                               y_pos=y_pos+height*0.5+1)
            # Display the sentences
            [visual.TextStim(self.window, text=sentence, pos=positions[s], color='black', height=text_height, wrapWidth=wrapWidth).draw() for s, sentence in enumerate(sentences)]
            self.show_presses(sentences, positions, pressed_keys, digit_start_time, wrapWidth, text_height)
            self.window.flip()

            if num_presses < num_items:
                keys = event.getKeys(keyList=self.const.response_keys, timeStamped=self.ttl_clock.clock)
                if keys:
                    key_char, key_press_time = keys[0]
                    key = self.const.response_keys.index(key_char) + 1
                    rt = key_press_time - digit_start_time
                    rt_list[num_presses]=rt
                    digit_start_time = key_press_time

                    # Check if key pressed is correct
                    correct_list[num_presses] = key == int(sequence[num_presses])
                    num_presses += 1
                    pressed_keys.append(key)

        # if any press is wrong trial['correct'] needs to be false, this is for post trial feedback
        trial['correct'] = correct_list.sum()/num_items

        if np.all(np.isnan(rt_list)):
            # calculate mean rt across presses
            trial['rt'] = np.nan

        else:
            trial['rt'] = np.nanmean(rt_list)

        # display trial feedback (for whole trial)
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct']==1)

        return trial

class ActionPrediction(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0],self.trial_info['key_two'].iloc[0]]

    def display_instructions(self):
        """
        displays the instruction for the task
        """
        task_name = visual.TextStim(self.window, text=f'{self.descriptive_name.capitalize()}', height=self.const.instruction_text_height, color=[-1, -1, -1], bold=True, pos=(0, 3))
        task_name.draw()
        self.instruction_text = ""
        if 'soccer' in self.task_file:
            self.instruction_text += "\n\n Decide if the ball is going to the left or right."
            self.instruction_text += f"\n\n\nLEFT: index finger \tRIGHT: middle finger\n"
        elif 'greeting' in self.task_file:
            self.instruction_text += "\n\n Decide if the people will hug or shake hands."
            self.instruction_text += f"\n\n\nHUG: index finger \tSHAKE HANDS: middle finger\n"
        else:
            self.instruction_text += "\n\n Choose what will happen next." # General instruction for both soccer and greeting
            self.instruction_text += f"\n\n\nLEFT/HUG: index finger \n\tRIGHT/SHAKE HANDS: middle finger\n"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1], wrapWidth=25, pos=(0, 0))
        instr_visual.draw()
        self.window.flip()


    def run_trial(self, trial):
        """ Runs a single trial of the Action Prediction task """

        event.clearEvents()

        window_width, _ = self.window.size
        movie_scale = getattr(self.const, 'action_prediction_scale', None) or 0.4
        stim_width = int(window_width * movie_scale) # Make the video fraction of the window width
        stim_height = int(stim_width  * 476 / 846)  # Original size of the video is 640x360


        # Display video
        movie_path = Path(self.const.stim_dir) / self.name / 'clips' / f"{trial['stim']}.mp4"
        movie_path_str = str(movie_path)
        movie_clip = visual.MovieStim(self.window, movie_path_str, loop=False, noAudio=True, size=(stim_width, stim_height), pos=(0, 0))

        movie_clip.play()

        self.window.flip()


        while movie_clip.isFinished == False:
            movie_clip.play()
            movie_clip.draw()
            self.window.flip()
            self.ttl_clock.update()
            # core.wait(1)  # Freeze the video for a moment

        # Display question
        options = trial['options'].split(',')
        question = trial['question']
        question += f"\n\n\n{self.corr_key[0]}. {options[0]} \t\t\t{self.corr_key[1]}. {options[1]}"
        question_stim = visual.TextStim(self.window, text=question, pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height= 1.5, wrapWidth=25)
        question_stim.draw()
        self.window.flip()

        # collect responses 0: no response 1-4: key pressed
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['question_dur'])
        trial['correct'] = (trial['response'] == options.index(str(trial['answer']))+1)

        # display trial feedback
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])

        # Flush memory: This is necessary for the script to be able to run more than 1 run. Presenting movies is very memory hungry, so do not remove!
        movie_clip.unload()
        gc.collect() # Collect garbarge

        return trial

class Movie(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.name = 'movie'

    def display_instructions(self):
        task_name = visual.TextStim(self.window, text=f'{self.descriptive_name.capitalize()}', height=self.const.instruction_text_height, color=[-1, -1, -1], bold=True, pos=(0, 3))
        task_name.draw()

        self.instruction_text = f"\n\n You will watch short clips from a movie. Please keep your head still and pay attention to the screen."
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1], wrapWidth=20, pos=(0, 0))
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        window_width, _ = self.window.size
        movie_scale = getattr(self.const, 'movie_scale', None) or 0.4
        stim_width = int(window_width * movie_scale) # Make the video fraction of the window width
        stim_height = int(stim_width  * 360 / 640)  # Original size of the video is 640x360

        # Get the file name
        movie_file_name = trial['stim']

        # Construct the movie file path
        movie_path = Path(self.const.stim_dir) / self.name / 'clips' / movie_file_name

        # Convert Path object to string for compatibility
        movie_path_str = str(movie_path)

        # Create a MovieStim3 object
        movie_clip = visual.MovieStim(self.window, movie_path_str, loop=False, size=(stim_width, stim_height), pos=(0, 0), noAudio=True)


        movie_clip.play()
        self.window.flip()

        while movie_clip.isFinished == False:
            movie_clip.play()
            movie_clip.draw()
            self.window.flip()
            self.ttl_clock.update()

        # Flush memory: This is necessary for the script to be able to run more than 1 run. Presenting movies is very memory hungry, so do not remove!
        movie_clip.unload()
        gc.collect() # Collect garbarge

        return trial


class StrangeStories(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.name = 'strange_stories'

    def init_task(self):
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0],self.trial_info['key_two'].iloc[0], self.trial_info['key_three'].iloc[0]]

    def display_instructions(self):
        task_name = visual.TextStim(self.window, text=f'{self.descriptive_name.capitalize()}', height=self.const.instruction_text_height, color=[-1, -1, -1], bold=True, pos=(0, 4))
        task_name.draw()

        self.instruction_text = f"\n\nYou will watch a clip about a couple and answer a question "
        if 'social' in self.task_file:
            self.instruction_text += "about their SOCIAL INTERACTION."
        elif 'control' in self.task_file:
            self.instruction_text += "about the FACTS."
        # self.instruction_text += " They live and work together. Each clip is self-contained and there is no story running from one clip to another."
        # self.instruction_text += "\n\n You will be asked a question about the clip. Imagine your answer as soon as you see the question. When you see the answer options, press the button that corresponds most to the answer you thought of. Some questions do not have a right or wrong answer."
        self.instruction_text += "\n\n Imagine your answer to the question. \nChoose the best match from the answers. \nSome questions have no right answer."
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1], wrapWidth=20, pos=(0, 0))
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        window_width, _ = self.window.size
        strange_stories_scale = getattr(self.const, 'strange_stories_scale', None) or 0.6
        stim_width = int(window_width * strange_stories_scale) # Make the video 40% of the window width
        stim_height = int(stim_width  * 921 / 1638)  # 1280x720 is the original size of the video given in width x height
        wrapWidth = 25

        # Get the file name
        movie_file_name = trial['stim']

        # Construct the movie file path
        movie_path = Path(self.const.stim_dir) / self.name / 'clips' / movie_file_name

        # Convert Path object to string for compatibility
        movie_path_str = str(movie_path)

        # Play the audio separately for better memory management
        play_audio_separatly = True

        # Create a MovieStim object
        movie_clip = visual.MovieStim(self.window, movie_path_str,
                                      loop=False, size=(stim_width, stim_height),
                                      pos=(0, 0), noAudio=play_audio_separatly)

        if play_audio_separatly:
            audio = self.get_audio_from_movie(movie_path, sample_rate=48000)

        movie_clip.draw()
        if play_audio_separatly:
            audio.play()
        movie_clip.play()
        self.window.flip()

        # Play through the movie frame by frame
        while movie_clip.isFinished == False:
            movie_clip.play()
            movie_clip.draw()
            self.window.flip()
            self.ttl_clock.update()

        if play_audio_separatly:
            audio.stop()

        # Initialize question
        question = trial['question']

        # Initialize answer options
        options_orig = [answer_option.strip() for answer_option in trial['options'].split(',')]
        options_shuffled = deepcopy(options_orig)
        random.shuffle(options_shuffled) # Randomize the order of the answer options
        if 'control' in trial['condition']: # Only the first option is correct (2 points)
            scores_orig = [2,0,0]
        elif 'social'in trial['condition']: # First option gets 2 points, second option gets 1 point, third option gets 0 points
            scores_orig = [2,1,0]
        scores_shuffled = [scores_orig[options_orig.index(option)] for option in options_shuffled]

        answers = f"\n\n\n{self.corr_key[0]}. {options_shuffled[0]} \n{self.corr_key[1]}. {options_shuffled[1]} \n{self.corr_key[2]}. {options_shuffled[2]}"

        # Display question
        stim_question = visual.TextStim(self.window, text = question, pos=(0, 4), color=(-1, -1, -1), units='deg', height= 1.5, wrapWidth=wrapWidth)
        stim_question.draw()
        self.window.flip()

        # Display the question until X seconds before trial is over (answer_dur), to make the 'buffer' zone for the trial, i.e. the time of variable length, the time where the participant deliberates about their answer
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + (trial['trial_dur'] - movie_clip.duration - trial['answer_dur']))
        # Flush any keys in buffer
        event.clearEvents()

        # Align the answers with the middle of the question if the answers are shorter than half of the question
        answer_lengths = [len(answer) for answer in options_shuffled]
        if max(answer_lengths) < wrapWidth and max(answer_lengths) < len(question):
            left_position = 0-max(answer_lengths)/3  # Answer options are shorter than questions and shorter than wrapWidth
            align='left'
        elif max(answer_lengths) >= wrapWidth:
            left_position = 0
            align='center'
        else:
            left_position = 0
            align='center'


        stim_answers = visual.TextStim(self.window, text=answers, pos=(left_position, 0), color=(-1, -1, -1), units='deg', height= 1.5, wrapWidth=wrapWidth, alignHoriz=align)
        stim_question.draw()
        stim_answers.draw()
        self.window.flip()

        # collect responses 0: no response 1-4: key pressed
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['answer_dur'])

        # Get the score for the selected answer
        if trial['response'] > len(scores_shuffled):
            trial['acc'] = 0 # If the participant pressed a key that is not in the list of answers, set the score to 0
        else:
            trial['acc'] = scores_shuffled[trial['response']-1]


        # Flush memory: This is necessary for the script to be able to run more than 1 run. Presenting movies is very memory hungry, so do not remove!
        movie_clip.unload()
        gc.collect() # Collect garbarge

        return trial


class FauxPas(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')
        self.corr_key = [self.trial_info['key_yes'].iloc[0],self.trial_info['key_no'].iloc[0]]


    def display_instructions(self):
        """
        displays the instruction for the task
        """
        task_name = visual.TextStim(self.window, text=f'{self.descriptive_name.capitalize()}', height=self.const.instruction_text_height, color=[-1, -1, -1], bold=True, pos=(0, 3))
        task_name.draw()
        self.instruction_text = "\n\nRead the story and answer the Yes/No question "
        if 'social' in self.task_file:
            self.instruction_text += "about the SOCIAL INTERACTION."
        elif 'control' in self.task_file:
            self.instruction_text += "about the FACTS."
        self.instruction_text += f"\n\n{self.corr_key[0]}. Yes \t{self.corr_key[1]}. No\n"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1], wrapWidth=20, pos=(0, 0))
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Runs a single trial of the Theory of Mind task """

        event.clearEvents()

        height = trial.get('text_height', 1.25)
        # Display story
        story = trial['story']
        # story = '.\n'.join(story.split('. '))
        story_stim = visual.TextStim(self.window, text=story, alignHoriz='center', wrapWidth=20, pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height= height)
        story_stim.draw()
        self.window.flip()

        # wait until story duration
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + trial['story_dur'])

        # Flush any keys in buffer
        event.clearEvents()

        # Display question
        question = trial['question']
        # Display answers
        options = [option.strip(' ') for option in trial['options'].split(',')]
        question += f"\n\n\n{self.corr_key[0]}. {options[0]} \t\t\t{self.corr_key[1]}. {options[1]}"
        question_stim = visual.TextStim(self.window, text=question, pos=(0.0, 0.0), color=(-1, -1, -1), units='deg', height= height, wrapWidth=25)
        question_stim.draw()
        self.window.flip()

        # collect responses 0: no response 1-4: key pressed
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(),
                                                           trial['question_dur'],
                                                           show_last_seconds=3,
                                                           current_stimuli=question_stim)
        trial['correct'] = (trial['response'] == trial['trial_type'])


        # display trial feedback
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])

        return trial


class FrithHappe(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.name = 'frith_happe'
        self.feedback_type = 'acc+rt'

    def init_task(self):
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0],self.trial_info['key_two'].iloc[0], self.trial_info['key_three'].iloc[0]]

    def display_instructions(self):
        task_name = visual.TextStim(self.window, text=f'{self.descriptive_name.capitalize()}', height=self.const.instruction_text_height, color=[-1, -1, -1], bold=True, pos=(0, 4))
        task_name.draw()

        self.instruction_text = f"Decide how the two triangles are interacting."
        instr_stim = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1], wrapWidth=20, pos=(0, 3))
        instr_stim.draw()
        # answer_expalantion = f"\n\n{self.corr_key[0]}. No interaction\n\n{self.corr_key[1]}. Physical (The actions are directed towards each other) \n\n{self.corr_key[2]}. Mental (One triangle manipulates the thoughts or feelings of the other)"
        answer_expalantion = f"\n\n{self.corr_key[0]}. No interaction\n\n{self.corr_key[1]}. Mental (One triangle manipulates the thoughts or feelings of the other)"
        instr_visual = visual.TextStim(self.window, text=answer_expalantion, color=[-1, -1, -1], wrapWidth=20, pos=(-8, -1), alignHoriz='left')
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Runs a single trial of the Frith-Happe task """
        # Flush any keys in buffer
        event.clearEvents()

        window_width, _ = self.window.size
        frith_happe_scale = getattr(self.const, 'frith_happe_scale', None) or 0.4
        stim_width = int(window_width * frith_happe_scale)
        stim_height = int(stim_width  * 1074 / 1433)
        wrapWidth = 25

        # Get the file name
        movie_file_name = trial['stim']
        # Construct the movie file path
        movie_path = Path(self.const.stim_dir) / self.name / 'clips' / movie_file_name
        # Convert Path object to string for compatibility
        movie_path_str = str(movie_path)
        # Create a MovieStim object
        movie_clip = visual.MovieStim(self.window, movie_path_str, loop=False, size=(stim_width, stim_height), pos=(0, 0), noAudio=True)


        movie_clip.play()
        self.window.flip()

        # Play through the movie frame by frame
        while movie_clip.isFinished == False:
            movie_clip.play()
            movie_clip.draw()
            self.window.flip()
            self.ttl_clock.update()

        # Initialize question
        question = "What type of interaction did you see?"
        # Display question
        stim_question = visual.TextStim(self.window, text = question, pos=(0, 2), color=(-1, -1, -1), units='deg', height= 1.5, wrapWidth=wrapWidth)
        stim_question.draw()
        self.window.flip()

        # Display the question until X seconds before trial is over (answer_dur), to make the 'buffer' zone for the trial, i.e. the time of variable length, the time where the participant deliberates about their answer
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + (trial['trial_dur'] - movie_clip.duration - trial['question_dur']))

        # Flush any keys in buffer
        event.clearEvents()

        stim_question.draw()
        # Initialize answer options
        # answers = f"\n\n{self.corr_key[0]}. No interaction \n{self.corr_key[1]}. Mental \n{self.corr_key[2]}. Physical "
        answers = f"\n\n{self.corr_key[0]}. No interaction \n{self.corr_key[1]}. Mental"
        answers_stim = visual.TextStim(self.window, text=answers, pos=(-5, 0), color=(-1, -1, -1), units='deg', height= 1.5, wrapWidth=wrapWidth, alignHoriz='left')
        answers_stim.draw()
        self.window.flip()

        # collect responses 0: no response 1-4: key pressed
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['question_dur'])
        trial['correct'] = (trial['response'] == trial['trial_type'])

        # display trial feedback
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])

        # Flush memory: This is necessary for the script to be able to run more than 1 run. Presenting movies is very memory hungry, so do not remove!
        movie_clip.unload()
        gc.collect() # Collect garbarge

        return trial



class Liking(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.name = 'liking'
        self.feedback_type = 'acc+rt'

    def init_task(self):
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0],self.trial_info['key_two'].iloc[0]]

        # MINIMAL: fallback if CSV has no key_one/key_two
        if {'key_one', 'key_two'}.issubset(self.trial_info.columns):
            self.corr_key = [str(self.trial_info['key_one'].iloc[0]), str(self.trial_info['key_two'].iloc[0])]
        else:
            # use constants (e.g., ['1','2','3','4']) → take first two
            fallback = getattr(self.const, 'response_keys', ['1', '2'])[:2]
            self.corr_key = [str(fallback[0]), str(fallback[1])]

    def display_instructions(self):
        task_name = visual.TextStim(self.window, text=f'{self.descriptive_name.capitalize()}',
                                    height=self.const.instruction_text_height, color=[-1, -1, -1],
                                    bold=True, pos=(0, 3))
        task_name.draw()

        self.instruction_text = "You will watch two people meeting for the first time.\n"

        # MINIMAL: normalize filename and set a safe default for key_text
        fname = str(self.task_file).lower().strip()
        key_text = f"\n{self.corr_key[0]}. Option 1 \t{self.corr_key[1]}. Option 2"

        if 'like' in fname and 'control' not in fname:
            self.instruction_text += "Judge if they LIKE each other."
            key_text = f"\n{self.corr_key[0]}. Yes \t{self.corr_key[1]}. No"
        elif 'control' in fname and 'like' not in fname:
            self.instruction_text += "Judge if one person SPEAKS MORE."
            key_text = f"\n{self.corr_key[0]}. Yes \t{self.corr_key[1]}. No"
        elif 'like' in fname and 'control' in fname:
            self.instruction_text += "Judge if they LIKE each other or if one person SPEAKS MORE."
            key_text = f"\n{self.corr_key[0]}. Yes \t{self.corr_key[1]}. No"
        # else: keep the default instruction_text + key_text

        instr_stim = visual.TextStim(self.window, text=self.instruction_text,
                                     height=self.const.instruction_text_height, color=[-1, -1, -1],
                                     wrapWidth=20, pos=(0, 0))
        instr_stim.draw()

        key_text = visual.TextStim(self.window, text=key_text,
                                   height=self.const.instruction_text_height, color=[-1, -1, -1],
                                   wrapWidth=20, pos=(-3, -3), alignHoriz='left')
        key_text.draw()
        self.window.flip()

    def run_trial(self, trial):
        window_width, _ = self.window.size
        liking_scale = getattr(self.const, 'liking_scale', None) or 0.5
        stim_width = int(window_width * liking_scale)
        stim_height = int(stim_width  * 486 / 720)
        wrapWidth = 20

        # Get the file name
        movie_file_name = trial['stim']
        # Construct the movie file path
        movie_path = Path(self.const.stim_dir) / self.name / 'clips' / (movie_file_name + '.mov')
        # Convert Path object to string for compatibility
        movie_path_str = str(movie_path)

        play_audio_separatly = True
        if play_audio_separatly:
            # Play the audio from the movie
            audio = self.get_audio_from_movie(movie_path, sample_rate=48000)

        # Create a MovieStim object
        movie_clip = visual.MovieStim(self.window, movie_path_str, loop=False,
                                      size=(stim_width, stim_height),
                                      pos=(0, 0),noAudio=play_audio_separatly)

        # Play through the movie frame by frame
        max_video_duration = 24
        movie_start_time = self.ttl_clock.get_time()

        movie_clip.draw()
        if play_audio_separatly:
            audio.play()
        movie_clip.play()
        self.window.flip()

        while (self.ttl_clock.get_time() - movie_start_time < max_video_duration):
            movie_clip.play()
            movie_clip.draw()
            self.window.flip()
            self.ttl_clock.update()

        # Flush any keys in buffer
        event.clearEvents()

        if play_audio_separatly:
            audio.stop()

        # Initialize question
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['question_dur'])
        if 'like' in trial['condition']:
            question = "Do they like each other?"
        elif 'control' in trial['condition']:
            question = "Did one person talk more?"

        # Display question
        stim_question = visual.TextStim(self.window, text = question, pos=(0, 1), color=(-1, -1, -1), units='deg', height= 1.5, wrapWidth=wrapWidth)
        stim_question.draw()

        # Initialize answer options
        answers = f"\n\n{self.corr_key[0]}. Yes \t{self.corr_key[1]}. No"
        stim_answers = visual.TextStim(self.window, text=answers, pos=(-3, 0), color=(-1, -1, -1), units='deg', height= 1.5, wrapWidth=wrapWidth, alignHoriz='left')
        stim_answers.draw()
        self.window.flip()

        # collect responses 0: no response 1-4: key pressed
        trial['response'],trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['question_dur'])
        if 'like' in trial['condition']:
            if trial['answer'] == 'like':
                trial['correct'] = (trial['response'] == 1)
            elif trial['answer'] == 'dislike':
                trial['correct'] = (trial['response'] == 2)
        elif 'control' in trial['condition']:
            if trial['answer'] == 'unbalanced':
                trial['correct'] = (trial['response'] == 1)
            elif trial['answer'] == 'balanced':
                trial['correct'] = (trial['response'] == 2)
        else:
            trial['correct'] = False

        # Record the played video duration
        trial['video_dur_orig'] = trial['video_dur']
        trial['video_dur'] = max_video_duration

        # display trial feedback
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])

        # Flush memory: This is necessary for the script to be able to run more than 1 run. Presenting movies is very memory hungry, so do not remove!
        movie_clip.unload()
        gc.collect() # Collect garbarge

        return trial

class Pong(Task):
    """
    Pong task
    """
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.name = 'pong'
        self.feedback_type = 'acc'

    def init_task(self):
        # Read trial info and set keys from file
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')
        self.corr_key = [self.trial_info['key_left'].iloc[0],
                         self.trial_info['key_right'].iloc[0]]
        # for real time handle movment
        self.key_handler = key.KeyStateHandler()
        self.window.winHandle.push_handlers(self.key_handler)

    def display_instructions(self):
        self.instruction_text = f"Use the buttons to move the paddle and catch the ball."
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, color=[-1, -1, -1],pos=(0, 0.3))
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        # Set parameters (all values are in degrees)
        paddle_speed = 0.5        # Movement per frame (deg)
        paddle_width = 7         # Paddle width (deg)
        paddle_height = 0.6        # Paddle thickness (deg)
        ball_radius = 0.8          # Ball radius (deg)
        trial_duration = trial['trial_dur']

        # Compute the effective screen dimensions (in degrees) based on monitor calibration.
        # Get monitor width (in cm) and viewing distance (in cm) from your screen object.
        monitor_width_cm = self.screen.monitor.getWidth()
        distance_cm = self.screen.distance

        # Calculate horizontal visual angle (in degrees)
        half_width_deg = math.degrees(math.atan((monitor_width_cm / 2) / distance_cm))
        screen_width_deg = 2 * half_width_deg

        # Compute vertical visual angle using the aspect ratio from the pixel dimensions.
        aspect_ratio = self.screen.size[1] / self.screen.size[0]
        screen_height_deg = screen_width_deg * aspect_ratio

        half_screen_width = screen_width_deg / 2.0
        half_screen_height = screen_height_deg / 2.0

        # Compute paddle boundaries in degrees
        paddle_half_width = paddle_width / 2.0
        min_x = -half_screen_width + paddle_half_width
        max_x = half_screen_width - paddle_half_width

        # Define margins (in degrees)
        paddle_margin = 2.0
        ball_margin = 2.0

        # Clear events
        event.clearEvents()

        # Convert the trajectory string to floats
        stim_str = trial['stim'].strip("()")
        dx, dy = map(float, stim_str.split(","))

        # Set initial positions for paddle and ball (in degrees)
        paddle_y = -half_screen_height + paddle_margin
        ball_y = half_screen_height - ball_margin

        # Create the paddle and ball stimuli using self.window (coordinates in deg)
        paddle = visual.Rect(self.window, width=paddle_width, height=paddle_height,
                             fillColor="white", pos=(0, paddle_y))
        ball = visual.Circle(self.window, radius=ball_radius, fillColor="white", pos=(0, ball_y))

        start_time = self.ttl_clock.get_time()

        # get movement keys
        key_left = getattr(key, self.const.response_keys[self.corr_key[0]-1].upper(), None)
        key_right = getattr(key, self.const.response_keys[self.corr_key[1]-1].upper(), None)
        trial['correct'] = False


        ball_stuck = False
        ball_offset_x = 0

        while self.ttl_clock.get_time() - start_time < trial_duration:
            if self.key_handler[key_left]:
                paddle.pos = (paddle.pos[0] - paddle_speed, paddle.pos[1])
            if self.key_handler[key_right]:
                paddle.pos = (paddle.pos[0] + paddle_speed, paddle.pos[1])

            # Clamp paddle position so it doesn't go off-screen
            if paddle.pos[0] < min_x:
                paddle.pos = (min_x, paddle.pos[1])
            if paddle.pos[0] > max_x:
                paddle.pos = (max_x, paddle.pos[1])

            # Update the ball position
            if not ball_stuck:
                ball.pos = (ball.pos[0] + dx, ball.pos[1] + dy)
            else:
                ball.pos = (paddle.pos[0] + ball_offset_x, paddle_y + paddle_height + ball_radius)


            # Bounce the ball off the side walls
            if ball.pos[0] >= half_screen_width - ball_radius or ball.pos[0] <= -half_screen_width + ball_radius:
                dx *= -1


            # Stick the ball to the paddle if it hits
            if not ball_stuck and dy < 0 and (paddle_y - ball_radius) < ball.pos[1] < (paddle_y + paddle_height + ball_radius):
                if (paddle.pos[0] - paddle_half_width) < ball.pos[0] < (paddle.pos[0] + paddle_half_width):
                    dy = 0
                    dx = 0
                    ball_stuck = True
                    ball_offset_x = ball.pos[0] - paddle.pos[0]  # Remember how far from center it landed
                    trial['correct'] = True

            # Draw the stimuli and update the display
            ball.draw()
            paddle.draw()
            self.window.flip()

        # Provide trial feedback
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])
        return trial

class Affective(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        """
        Initialize task
        """
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')

        self.stim = []
        for stim_file in self.trial_info['stim']:
            stim_path = self.const.stim_dir / self.name / stim_file
            self.stim.append(visual.ImageStim(self.window, image=str(stim_path)))

    def display_instructions(self):
        """
        Display instructions for the affective task.
        """
        key_pleasant = self.trial_info['key_pleasant'].iloc[0]
        key_unpleasant = self.trial_info['key_unpleasant'].iloc[0]

        self.instruction_text = (
            f"{self.descriptive_name} Task\n\n"
            f"Press {key_unpleasant} if the image is UNPLEASANT.\n"
            f"Press {key_pleasant} if the image is PLEASANT."
        )
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1], wrapWidth=20)
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """
        Run a single trial of the affective task.
        """
        event.clearEvents()

        # Display image
        self.stim[trial['trial_num']].draw()
        self.window.flip()

        # Wait for response
        trial['response'], trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['trial_dur'])

        # Determine correct response
        correct_key = trial['key_pleasant'] if trial['trial_type'] == 2 else trial['key_unpleasant']
        trial['correct'] = (trial['response'] == correct_key)

        # Show feedback
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])

        return trial

class SerialReactionTime(Task):
    """
    Serial reaction time task
    """
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)

    def init_task(self):
        """
        Initialize task - default is to read the target information into the trial_info dataframe
        """
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0],self.trial_info['key_two'].iloc[0],self.trial_info['key_three'].iloc[0],self.trial_info['key_four'].iloc[0]]


    def display_instructions(self): # overriding the display instruction from the parent class
        self.instruction_text = f"Press the buttons that match the green boxes"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, color=[-1, -1, -1],pos=(0, 0.3))
        instr_visual.draw()
        self.window.flip()

    def run(self):
        """Override to show empty boxes during initial wait."""
        box_positions = [(-5, 2), (-3, 2), (3, 2), (5, 2)]
        self.boxes = [
            visual.Rect(
                self.window, width=2, height=2,
                pos=pos, fillColor='white',
                lineColor='black', units='deg'
            )
            for pos in box_positions
        ]
        for box in self.boxes:
            box.draw()
        self.window.flip()
        return super().run()

    def run_trial(self, trial):

        event.clearEvents()

        trial_stim = int(trial['stim'])
        target_index = trial_stim - 1

        # wait until absolute stimulus time
        stim_onset = self.ttl_clock.get_time()

        # show stimulus
        for j, box in enumerate(self.boxes):
            box.fillColor = 'green' if j == target_index else 'white'
            box.draw()
        self.window.flip()

        # keep green on for trial_dur
        self.ttl_clock.wait_until(stim_onset + trial['trial_dur'])

        # back to white
        for box in self.boxes:
            box.fillColor = 'white'
            box.draw()
        self.window.flip()

        # collect responses until absolute end_time
        responses = []
        while self.ttl_clock.get_time() < (stim_onset + trial['trial_dur'] + trial['iti_dur']):
            keys = event.getKeys(
                keyList=self.const.response_keys,
                timeStamped=self.ttl_clock.clock
            )
            for key, timestamp in keys:
                rt = timestamp - stim_onset
                responses.append((key, rt))

        trial['responses'] = responses
        return trial
    
class FingerRhythmic(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)

    def init_task(self):
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0]]

    def display_instructions(self):
        """
        displays the instruction for the task
        """

        str1 = f"Tap along to the tones using the {self.corr_key[0]} key."
        str2 = f"Keep tapping at the same pace when the tones stop."
        self.instruction_text = f"{self.descriptive_name} Task\n\n {str1} \n {str2}"
        instr_visual = visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1])
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):
        """ Runs a single trial of the Finger Rhythmic task """

        event.clearEvents()
        txt = (f"New trial starts now") # this text shows when a new trials starts
        visual.TextStim(self.window, text=txt,height=self.const.instruction_text_height, color=[-1, -1, -1]).draw()
        self.window.flip()
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + 2)

        self.screen.fixation_cross()
        event.clearEvents()
        clk = self.ttl_clock.clock
        t0 = clk.getTime()                    # trial anchor (TTL)

        # --- Play FIRST tone now, then use THIS time as the grid anchor
        beep = sound.Sound(value=1000, secs=0.05, sampleRate=48000, stereo=True)
        beep.play()
        t_first = clk.getTime()               # when we triggered the first tone (TTL)
        ioi = 0.65
        expected = [(t_first - t0) + i*ioi for i in range(12)]  # expected, aligned to first tone

        # Track beep times for timing verification (quiet by default)
        beep_times = [t_first]

        taps_rel = []

        # --- Remaining 11 tones by absolute TTL deadlines; collect keys in-between
        for i in range(1, 12):
            deadline = t_first + i*ioi
            while True:
                now = clk.getTime()
                if now >= deadline:
                    # Create a new Sound object for each beep to ensure it plays
                    beep = sound.Sound(value=1000, secs=0.05, sampleRate=48000, stereo=True)
                    beep.play()
                    beep_time = clk.getTime()
                    beep_times.append(beep_time)
                    break
                res = event.waitKeys(maxWait=deadline - now,
                                     keyList=self.const.response_keys,
                                     timeStamped=clk)
                if res:
                    for _, ts in res:
                        taps_rel.append(ts - t0)

        # --- Silent phase: collect until absolute end_time
        end_abs = float(trial['end_time'])
        while True:
            now = clk.getTime()
            if now >= end_abs:
                break
            res = event.waitKeys(maxWait=end_abs - now,
                                 keyList=self.const.response_keys,
                                 timeStamped=clk)
            if res:
                for _, ts in res:
                    taps_rel.append(ts - t0)

        # --- Self-paced ISIs only (strictly after last tone onset)
        last_tone_t = expected[-1]            # relative to t0
        self_taps = [t for t in taps_rel if t > last_tone_t]
        isis = np.diff(self_taps) if len(self_taps) > 1 else np.array([], float)
        isis = isis[(isis >= 0.300) & (isis <= 0.900)]

        trial['iri_ms_mean']         = float(np.mean(isis) * 1000.0) if isis.size else np.nan
        trial['iri_ms_sd']           = float(np.std(isis)  * 1000.0) if isis.size else np.nan
        trial['iris_ms_json']        = json.dumps((isis * 1000.0).tolist())
        trial['expected_rel_s_json'] = json.dumps([e for e in expected])  # seconds rel to t0
        trial['tap_rel_s_json']      = json.dumps(taps_rel)
        # Save beep times relative to trial start (t0)
        beep_times_rel = [bt - t0 for bt in beep_times]
        trial['beep_times_rel_s_json'] = json.dumps(beep_times_rel)

        return trial


class TimePerception(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        self.trial_info = pd.read_csv(self.const.task_dir / self.name / self.task_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0], self.trial_info['key_two'].iloc[0]]
        # one tone for both modalities
        self.tone = sound.Sound(value=1000, secs=0.050, sampleRate=48000, stereo=True)  # 1000 Hz, 50 ms

        # PEST state per side
        mod = str(self.trial_info['modality'].iloc[0]).lower()
        if mod == 'time':
            # start ±120 ms from 400; step 40; min step 8; directions tracked
            self.stair = {
                'shorter': {'curr': 280.0, 'step': 40.0, 'min_step': 8.0,  'last_dir': 0, 'same_dir': 0},
                'longer':  {'curr': 520.0, 'step': 40.0, 'min_step': 8.0,  'last_dir': 0, 'same_dir': 0},
            }
        else:  # 'volume' (quieter/louder)
            # start ±1.62 dB from 73; step 1.08; min step 0.27
            self.stair = {
                'quieter': {'curr': 71.38, 'step': 1.08, 'min_step': 0.27, 'last_dir': 0, 'same_dir': 0},
                'louder':  {'curr': 74.62, 'step': 1.08, 'min_step': 0.27, 'last_dir': 0, 'same_dir': 0},
            }

    def display_instructions(self):
        mod = str(self.trial_info['modality'].iloc[0]).lower()
        if mod == 'time':
            str1 = f"You will hear two pairs of tones."
            str2 = f"Press [{self.corr_key[0]}] if the SECOND interval is shorter."
            str3 = f"Press [{self.corr_key[1]}] if the SECOND interval is longer."
            str4 = "The first pair is always the same."
            self.instruction_text = f"{self.descriptive_name} Task\n\n {str1} \n {str2} \n {str3} \n {str4}"

        else:
            str1 = f"You will hear two pairs of tones."
            str2 = f"Press [{self.corr_key[0]}] if the SECOND interval is quieter."
            str3 = f"Press [{self.corr_key[1]}] if the SECOND interval is louder."
            str4 = "The first pair is always the same."
            self.instruction_text = f"{self.descriptive_name} Task\n\n {str1} \n {str2} \n {str3} \n {str4}"
        visual.TextStim(self.window, text=self.instruction_text, height=self.const.instruction_text_height, color=[-1, -1, -1]).draw()
        self.window.flip()

    def run_trial(self, trial):
        event.clearEvents()
        clk  = self.ttl_clock
        mod  = str(trial['modality']).lower()
        side = str(trial['side']).lower()

        # fixation
        self.screen.fixation_cross()

        # --- Pair 1 (standard @ 0.7 amp, 400 ms gap) ---
        self.tone.setVolume(0.7)
        t1 = clk.get_time()
        self.tone.play()
        clk.wait_until(t1 + 0.050)
        clk.wait_until(t1 + 0.050 + 0.400)
        self.tone.play()
        clk.wait_until(t1 + 2*0.050 + 0.400)

        # inter-pair gap
        clk.wait_until(clk.get_time() + 1.000)

        # --- Pair 2 (comparison) ---
        st = self.stair[side]  # PEST state for this side

        if mod == 'time':
            # snap to 8 ms grid within side range
            if side == 'shorter':
                comp_ms = max(160, min(392, int(round(st['curr'] / 8.0) * 8)))
            else:
                comp_ms = max(408, min(640, int(round(st['curr'] / 8.0) * 8)))

            t2 = clk.get_time()
            self.tone.setVolume(0.7)
            self.tone.play()
            clk.wait_until(t2 + 0.050)
            clk.wait_until(t2 + 0.050 + (comp_ms / 1000.0))
            self.tone.play()
            clk.wait_until(t2 + 2*0.050 + (comp_ms / 1000.0))

            trial['comparison_ms'] = float(comp_ms)
            trial['stair_step_ms'] = float(st['step'])  # log step used this trial

        else:  # volume (quieter/louder), grid 0.27 dB
            if side == 'quieter':
                comp_db = max(64.9, min(72.73, float(st['curr'])))
            else:
                comp_db = min(81.1, max(73.27, float(st['curr'])))

            comp_amp = float(min(1.0, 0.7 * (10 ** ((comp_db - 73.0) / 20.0))))

            t2 = clk.get_time()
            self.tone.setVolume(comp_amp)
            self.tone.play()
            clk.wait_until(t2 + 0.050)
            clk.wait_until(t2 + 0.050 + 0.400)
            self.tone.setVolume(comp_amp)
            self.tone.play()
            clk.wait_until(t2 + 2*0.050 + 0.400)
            self.tone.setVolume(1.0)

            trial['comparison_dba']  = float(comp_db)
            trial['stair_step_dba']  = float(st['step'])  # log step used this trial

        # --- response window ---
        trial['response'], trial['rt'] = self.wait_response(clk.get_time(), float(trial['question_dur']))
        trial['correct'] = (trial['response'] == trial['trial_type'])

        # --- PEST update (classic + border-safe): halve on reversal; double after two same-direction moves ---
        # Define movement: toward standard if correct, away if incorrect.
        # Use unified sign for direction comparison: toward = -1, away = +1.
        move_dir = (-1 if trial['correct'] else +1)

        # Store the old level so we can tell if we actually moved after clamping.
        old_curr = st['curr']

        # 1. Propose + clamp
        if mod == 'time':
            if side == 'shorter':
                # correct => toward standard => make interval longer => +step
                # incorrect => away => make interval even shorter => -step
                st['curr'] += (+st['step'] if trial['correct'] else -st['step'])
                # snap to 8 ms grid and clamp to that side's allowed range
                st['curr']  = float(int(round(st['curr'] / 8.0) * 8))
                if st['curr'] < 160.0:
                    st['curr'] = 160.0
                if st['curr'] > 392.0:
                    st['curr'] = 392.0

            else:  # side == 'longer'
                # correct => toward standard => make interval shorter => -step
                # incorrect => away => make interval even longer => +step
                st['curr'] += (-st['step'] if trial['correct'] else +st['step'])
                st['curr']  = float(int(round(st['curr'] / 8.0) * 8))
                if st['curr'] < 408.0:
                    st['curr'] = 408.0
                if st['curr'] > 640.0:
                    st['curr'] = 640.0

        else:
            # volume
            if side == 'quieter':
                # correct => toward standard (louder) => +step in dB toward 73
                # incorrect => away (quieter) => -step
                st['curr'] += (+st['step'] if trial['correct'] else -st['step'])
                if st['curr'] < 64.9:
                    st['curr'] = 64.9
                if st['curr'] > 72.73:
                    st['curr'] = 72.73

            else:  # side == 'louder'
                # correct => toward standard (quieter) => -step
                # incorrect => away (louder) => +step
                st['curr'] += (-st['step'] if trial['correct'] else +st['step'])
                if st['curr'] < 73.27:
                    st['curr'] = 73.27
                if st['curr'] > 81.1:
                    st['curr'] = 81.1

        # 2. Did we actually move?
        actually_moved = (st['curr'] != old_curr)

        if actually_moved:
            # Normal PESt adaptation only if we escaped the boundary.

            if st['last_dir'] != 0 and move_dir != st['last_dir']:
                # reversal -> halve step (but not below min_step), reset consecutive counter
                st['step'] = max(st['min_step'], st['step'] / 2.0)
                st['same_dir'] = 0

            else:
                # same direction as last (or first informative move)
                if st['last_dir'] == 0 or move_dir == st['last_dir']:
                    st['same_dir'] += 1
                else:
                    # new direction but last_dir was 0 should already be covered above,
                    # but keep a safe fallback
                    st['same_dir'] = 1

                # after two same-direction moves -> double step
                if st['same_dir'] >= 2:
                    st['step'] = st['step'] * 2.0
                    st['same_dir'] = 0  # require two more same-direction moves for next doubling

            # update last_dir ONLY when a real move happened
            st['last_dir'] = move_dir

        else:
            # We hit a border and got clamped. Do NOT adapt step. Do NOT change same_dir. Do NOT touch last_dir.
            pass

        self.stair[side] = st
        self.display_trial_feedback(trial['display_trial_feedback'], trial['correct'])
        return trial

class SensMotControl(Task):
    def __init__(self, info, screen, ttl_clock, const, subj_id):
        super().__init__(info, screen, ttl_clock, const, subj_id)
        self.feedback_type = 'acc+rt'

    def init_task(self):
        trial_info_file = self.const.task_dir / self.name / self.task_file
        self.trial_info = pd.read_csv(trial_info_file, sep='\t')
        self.corr_key = [self.trial_info['key_one'].iloc[0], self.trial_info['key_two'].iloc[0]]

    def display_instructions(self):
        """
        displays the instruction for the task
        """
        cond = str(self.trial_info['condition'].iloc[0])
        if cond == 'blue':
            str1 = f"When the circle turns BLUE, press {self.corr_key[0]}."
            str2 = f"When the circle turns WHITE, do nothing."
            self.instruction_text = f"{self.descriptive_name} Task\n\n {str1} \n {str2}"
        elif cond == 'red':
            str1 = f"When the circle turns RED, press {self.corr_key[1]}."
            str2 = f"When the circle turns WHITE, do nothing."
            self.instruction_text = f"{self.descriptive_name} Task\n\n {str1} \n {str2}"
        else:
            str1 = f"When the circle turns BLUE, press {self.corr_key[0]}."
            str2 = f"When the circle turns RED, press {self.corr_key[1]}."
            str3 = f"When the circle turns WHITE, do nothing."
            self.instruction_text = f"{self.descriptive_name} Task\n\n {str1} \n {str2} \n {str3}"

        instr_visual = visual.TextStim(self.window, text=self.instruction_text,
                                       height=self.const.instruction_text_height, color=[-1, -1, -1], wrapWidth=25, pos=(0, 0))
        instr_visual.draw()
        self.window.flip()

    def run_trial(self, trial):

        event.clearEvents()

        # --- 1) Fixation circle (1000 ms) ---
        visual.Circle(self.window, radius=3, edges=128, lineWidth=4, lineColor='black', fillColor=None).draw()
        self.window.flip()
        self.ttl_clock.wait_until(self.ttl_clock.get_time() + 1)
        self.ttl_clock.update()

        # --- 2) Colored circle (2000 ms) ---
        visual.Circle(self.window, radius=3, edges=128,lineWidth=6, fillColor= trial['stim'], lineColor=trial['stim']).draw()
        self.window.flip()

        # collect responses 0: no response 1-4: key pressed
        trial['response'], trial['rt'] = self.wait_response(self.ttl_clock.get_time(), trial['question_dur'])
        trial['correct'] = (trial['response'] == trial['trial_type'])

        # display trial feedback
        if trial['display_trial_feedback']:
            # show feedback, then let the schedule absorb any remaining time
            self.display_trial_feedback(True, trial['correct'])
        else:
            # no feedback: go BLANK immediately and stay blank until end_time (we don't want the fixation cross here)
            self.window.flip(clearBuffer=True)
            while self.ttl_clock.get_time() < trial['end_time']:
                # flip blank frames so the window stays responsive
                self.window.flip()
                self.ttl_clock.update()

        return trial