import os
import platform
import random
from datetime import datetime
import logging
from psychopy import visual, core, event
from psychopy.gui import DlgFromDict
from psychopy.hardware.keyboard import Keyboard
from sequences import stimuli_manager as sm
from sequences import flow as fl
from sequences import params as pm
from sequences import instr as it
from sequences.common import get_win_dict, get_photo_rects
import byte_triggers as bt
from main import setup_sequence_distribution

def execute_loc():
    ''' Run the localizer. Returns nothing. '''
    tools = init_loc()
    logger = tools['logger']
    exp_info = tools['exp_info']
    loc_type = exp_info['loc_type']

    try:
        amodal_sequences, _ = setup_sequence_distribution(tools)
    except Exception as exc:
        fl.log_exceptions(f"An error occurred during sequence generation: {exc}", logger, tools['win'])
    
    # we get the names of sequences of run_01 of the crossmod memory task. This is used to select the stimuli for the localizer.
    two_run_org = sm.generate_run_org(seq_structures=pm.seq_structures, seed=tools['seed']) 
    run_01_org = two_run_org['run1'].values()
    run_01_seq = set([item for sublist in list(run_01_org) for item in sublist])  # Flatten the list of lists
    # extract the sequences and stims that are in run_01
    sequences = {seq_name:seq for seq_name, seq in amodal_sequences.items() if seq_name in run_01_seq}
    stimuli = {}
    for mod in ['img', 'txt']:
        for seq in sequences.values():
            stim_paths = sm.get_stims(pm.input_dir, seq, mod, lang=exp_info['lang'])
            stimuli.update({s_path: mod for s_path in stim_paths})
    
    print(f"N stimuli for localizer: {len(stimuli)}")

    # present all the blocks
    present_blocks( 
        tools=tools,
        stim_paths=stimuli,
        loc_type=loc_type,
    )
    end_txt = it.get_txt(exp_info['lang'], 'end_loc_fn')
    end_msg = visual.TextStim(
        win=tools['win'],
        text=end_txt,
        pos=(0, 0),
        font="Arial",
        color='black', 
        height=pm.text_height,
        units='norm'
    )

    # end message and close run
    tools['background'].draw() 
    end_msg.draw()  
    tools['win'].flip() 
    core.wait(2)
    fl.end_loc(tools)  # end the localizer run gracefully
    return

def init_loc():
    ''' Initialize the localizer run. Returns a dictionary with tools for the run. '''

    exp_info = {
        'ID': '00',
        'run': '01',
        'lang': 'fr',
        'loc_type': 'A', 
        'debugging': False
    }
    dlg = DlgFromDict(exp_info, title='Enter participant info', sortKeys=False)
    fl.handle_user_cancel(dlg)
    fl.check_user_info(exp_info)
    sm.check_img_txt(pm.input_dir, exp_info['lang'])
    out_dir = f"data/output/sub-{exp_info['ID']}"
    logfn = f"{out_dir}/sub-{exp_info['ID']}_loc-type-{exp_info['loc_type']}_run{exp_info['run']}-logs-{datetime.now().strftime('%Y%m%d-%H%M')}.log"
    os.makedirs(out_dir, exist_ok=True)

    if platform.system() == 'Linux':
        keyboard = Keyboard()
        keyboard.start()
        adapt_waitKeys = keyboard.waitKeys
        adapt_getKeys = keyboard.getKeys
        adapt_clearEvents = keyboard.clearEvents
    else:
        keyboard = None
        adapt_waitKeys = event.waitKeys
        adapt_getKeys = event.getKeys
        adapt_clearEvents = event.clearEvents

    bt.add_file_handler(
        logfn,
        mode='a', 
        verbose='INFO'
    )
    logging.basicConfig(
        filename=logfn,
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )

    if pm.use_mock_port:
        pport = bt.MockTrigger()
    else:
        pport = bt.ParallelPortTrigger(pm.pport, port_type='pport', delay=10)

    seed = sm.set_seed(exp_info['ID'])

    logger = logging.getLogger()
    logger.info(f'running on {platform.system()}')
    logger.info('Localizer started.')
    logger.info(f'Seed: {seed}')
    logger.info(f'Participant ID: {exp_info["ID"]}')
    logger.info(f'Run number: {exp_info["run"]}')
    logger.info(f'Language: {exp_info["lang"]}')

    win_dict = get_win_dict()
    win_dict['win'].mouseVisible = False

    tools = {
        'pport': pport,
        'logger': logger,
        'win': win_dict['win'],
        'aspect_ratio': win_dict['aspect_ratio'],
        'background': win_dict['background'],
        'wait_fun': core.wait,
        'event_fun': adapt_getKeys,
        'trig_fun': pport.signal,
        'clear_event_fun':adapt_clearEvents, 
        'exp_info': exp_info,
        'seed': seed,
        'out_dir': out_dir,
        'keyboard': keyboard,
        'adapt_waitKeys': adapt_waitKeys,
    }

    return tools


def present_blocks(tools, stim_paths, loc_type='A'):
    ''' Present a block of stimuli. Returns nothing. 
    Based on the loc_type, different stim oredering and different tasks to be performed. '''

    def get_resp_dict():
        ''' Initialize a response dictionary to keep track of responses. '''
        return {'correct': 0, 'false_alarm': 0, 'miss': 0}
    
    win = tools['win']
    background = tools['background']
    aspect_ratio = tools['aspect_ratio']
    lang = tools['exp_info']['lang']
    adapt_waitKeys = tools['adapt_waitKeys']
    logger = tools['logger']

    if loc_type == 'A':
        block_org_img, tgt_info_img = sm.prep_stim_tgt_order(
            input_dir=pm.input_dir,
            stim_paths=stim_paths, 
            modality='img', 
            nblocks=pm.loc_nblocks,
            n_tgt_low=pm.n_tgt_low,
            n_tgt_high=pm.n_tgt_high,
            lang=lang,
        )
        block_org_txt, tgt_info_txt = sm.prep_stim_tgt_order(
            input_dir=pm.input_dir,
            stim_paths=stim_paths, 
            modality='txt',
            nblocks=pm.loc_nblocks,
            n_tgt_low=pm.n_tgt_low,
            n_tgt_high=pm.n_tgt_high,
            lang=lang,
        )
        #logger.info(f'Block organization for img: {block_org_img}')
        #logger.info(f'Block organization for txt: {block_org_txt}')
        logger.info(f'Target info for img: {tgt_info_img}')
        logger.info(f'Target info for txt: {tgt_info_txt}')

    elif loc_type == 'B':
        block_org_img, supra_cat_order = sm.prep_stim_supra_cat(
            stim_paths, 
            modality='img',
            n_tgt_low=pm.n_tgt_low,
            n_tgt_high=pm.n_tgt_high,
        )
        block_org_txt, _ = sm.prep_stim_supra_cat( # should not be _ if randomized
            stim_paths, 
            modality='txt',
            n_tgt_low=pm.n_tgt_low,
            n_tgt_high=pm.n_tgt_high,
        ) 
        #logger.info(f'Block organization for img: {block_org_img}')
        #logger.info(f'Block organization for txt: {block_org_txt}')
        logger.info(f'Supra category order: {supra_cat_order}')

    elif loc_type == 'C':
        # in this case, there is no separation by modality. 
        block_org = sm.prep_stim_match(
            stim_paths,
            nblocks= pm.loc_nblocks,
            n_tgt_low=pm.n_tgt_low,
            n_tgt_high=pm.n_tgt_high,
        )
        logger.info(f'Block organization: {block_org}')

    # present instructions
    instr1_txt = it.get_txt(lang, f'instr_loc{loc_type}1_fn')
    instr1 = visual.TextStim(win=tools['win'],
        text=instr1_txt,
        pos=(0, 0),
        font="Arial",
        color='black', 
        height=pm.text_height,
        units='norm'
    )

    if loc_type != 'B':
        instr2_txt = it.get_txt(lang, f'instr_loc{loc_type}2_fn')
        instr2 = visual.TextStim(win=tools['win'],
            text=instr2_txt,
            pos=(0, 0.6),
            font="Arial",
            color='black', 
            height=pm.text_height,
            units='norm'
        )
    else:
        fns = ['instr_locB2_anim_fn', 'instr_locB2_inanim_fn']
        supra_cat = ['animate', 'inanimate']
        instr_supra = {}
        for fn, cat in zip(fns, supra_cat):
            instr_supra[cat] = visual.TextStim(win=tools['win'],
                text=it.get_txt(lang, fn),
                pos=(0, 0),
                font="Arial",
                color='black', 
                height=pm.text_height,
                units='norm'
            )

    background.draw()
    instr1.draw()   
    win.flip()
    adapt_waitKeys(keyList=[pm.key_dict['confirm']])
    
    # iterate over the blocks and present stimuli
    modalities = ['img', 'txt'] 
    for i in range(pm.loc_nblocks):
        print(f'======= Presenting block {i+1} of {pm.loc_nblocks} =======')
        logger.info(f'Presenting block {i+1} of {pm.loc_nblocks}')
        # actual stim presentation
        resp_dict = get_resp_dict()  # reset response dictionary for each block
        if loc_type == 'A':
            mod = modalities[i % 2] 
            # instructions for the block
            # alternate between img and txt
            if mod == 'img':
                stims = block_org_img[f'block{i+1}']
                tgt_stim_path = tgt_info_img[i]
            elif mod == 'txt':
                stims = block_org_txt[f'block{i+1}']
                tgt_stim_path = tgt_info_txt[i]
            logger.info(f'Modality for block {i+1}: {mod}')
            #logger.info(f'Stimuli for block {i+1} ({mod}): {stims}')
            #logger.info(f'Target stimulus for block {i+1} ({mod}): {tgt_stim_path}')

            tgt_stim = visual.ImageStim(
                win=win,
                image=tgt_stim_path,
                size=(pm.img_size, pm.img_size * aspect_ratio)
            )
            background.draw()
            tgt_stim.draw()
            instr2.draw()
            win.flip()
            adapt_waitKeys(keyList=[pm.key_dict['confirm']])

            for j, d in enumerate(stims):
                present_stimulus(tools, stim=d['path'], modality=mod, label=d['label'], resp_dict=resp_dict)
                print(j)
                fl.check_escape_or_break(tools, pause_key=pm.key_dict['pause'])
            give_feedback(tools, resp_dict, stims)

        elif loc_type == 'B':
            # don't alternate between modalities because it alternates between supra categories
            mod = 'txt' if i > 2 else 'img'
            # instructions for the block
            if mod == 'img':
                stims = block_org_img[f'block{i+1}']
            elif mod == 'txt':
                stims = block_org_txt[f'block{i+1}']

            supra_cat_name = supra_cat_order[i] 
            background.draw()
            instr_supra[supra_cat_name].draw()
            win.flip()
            adapt_waitKeys(keyList=[pm.key_dict['confirm']])

            for j, d in enumerate(stims):
                present_stimulus(tools, stim=d['path'], modality=mod, label=d['label'], resp_dict=resp_dict)
                print(j)
                fl.check_escape_or_break(tools, pause_key=pm.key_dict['pause'])
            give_feedback(tools, resp_dict, stims)

        elif loc_type == 'C':
            stims = block_org[f'block{i+1}']

            # instructions for the block
            background.draw()
            instr2.draw()
            win.flip()
            adapt_waitKeys(keyList=[pm.key_dict['confirm']])

            for j, d in enumerate(stims):
                mod = sm.get_stim_mod(d['path'])  # get modality from the path
                present_stimulus(tools, stim=d['path'], modality=mod, label=d['label'], resp_dict=resp_dict) # need to get the modality
                print(j)
                fl.check_escape_or_break(tools, pause_key=pm.key_dict['pause'])
            give_feedback(tools, resp_dict, stims)
        
        else:
            logger.error(f'Unknown localizer type: {loc_type}')
            print(f'Unknown localizer type: {loc_type}')
            fl.end_loc(tools)

    return

def give_feedback(tools, resp_dict, stims):
    ''' Give feedback on the responses. Returns nothing. '''
    win = tools['win']
    background = tools['background']
    lang = tools['exp_info']['lang']
    logger = tools['logger']

    # just get the number of target stimuli
    target_count = sum(1 for item in stims if item['label'] == 'target')

    # add block over message and feedback. 
    prop_incor = sm.cpt_points_loc(resp_dict)
    if prop_incor > pm.loc_incor_thresh or resp_dict['miss'] == target_count:
        feedback_txt = it.get_txt(lang, 'instr_loc_fb2_fn')
        color = 'black'
    else:
        feedback_txt = it.get_txt(lang, 'instr_loc_fb1_fn')
        color = 'green'

    feedback = visual.TextStim(win=win,
        text=feedback_txt.format(prop_incor),
        pos=(0, 0),
        font="Arial",
        color=color, 
        height=pm.text_height,
        units='norm'
    )
    logger.info('--- Feedback ---')
    logger.info(f'Number of targets in block: {target_count}')
    logger.info(f'Number of misses: {resp_dict["miss"]}')
    logger.info(f'Proportion of incorrect responses: {prop_incor}')
    logger.info(f'Feedback: {feedback_txt}')
    background.draw()
    feedback.draw()
    win.flip()
    core.wait(pm.t_loc_fb)

    return

def present_stimulus(tools, stim, modality, label, resp_dict): 
    ''' Present a single stimulus. Returns nothing. '''
    
    pport = tools['pport']
    logger = tools['logger']
    win = tools['win']
    aspect_ratio = tools['aspect_ratio']
    background = tools['background']
    debugging = tools['exp_info']['debugging']
    clock = core.Clock()

    t_stim = pm.loc_stim_dur 
    t_isi = pm.loc_isi_dur + random.uniform(0, pm.loc_jitter)
    if debugging:
        t_stim = 0.1
        t_isi = 0.1

    stim_cat = sm.get_cat_from_stim(stim)
    stim_name = sm.get_stim_name(stim)
    trig1 = pm.triggers['mod_cat'][modality][stim_cat] # keys are 'img'/'txt' and category names (e.g., 'animals')
    trig2 = pm.triggers['loc'][label] # trigger indicating the label of the stimulus (e.g., 'target'/'valid')
    trig3 = pm.triggers['misc']['confirm'] # used when participant presses the detection button
                        
    fl.check_escape_or_break(tools, pause_key=pm.key_dict['pause'])
    stim_image = visual.ImageStim(
        win=win,
        image=stim,
        size=(pm.img_size, pm.img_size*aspect_ratio)
    )

    # contrast for photodiode
    rect_w, rect_b = get_photo_rects(win, aspect_ratio)
    
    # Display fixation cross
    fix_cross = visual.TextStim(win=win,
        text='+',
        font='Arial',
        height=0.1,
        color='black',
        units='norm'
    )

    background.draw()
    rect_b.draw()
    rect_w.draw()
    stim_image.draw()
    # log info there to be closer to the actual presentation
    logger.info('--- Stim to be presented ---')
    logger.info(f'stimulus name: {stim_name}')
    logger.info(f'stimulus modality: {modality}')
    logger.info(f'stimulus category: {stim_cat}')
    logger.info(f'stimulus path: {stim}')
    logger.info(f'stimulus label: {label}')
    logger.info(f'trig1: {trig1}')
    logger.info(f'trig2: {trig2}')

    print(label)
    print(stim_name)
    win.callOnFlip(fl.novov_trigger,pport=pport, trig1=trig1, trig2=trig2, delay=10)
    win.flip()
    core.wait(t_stim)
    background.draw()
    rect_b.draw()
    fix_cross.draw()
    win.callOnFlip(clock.reset) 
    win.flip()
    resp = sm.wait_and_check_resp(tools, clock, t_isi, trig_val=trig3, key_dict=pm.key_dict)
    print(f'Button pressed: {resp}')
    logger.info(f'Button pressed: {resp}')
    sm.feed_resp_dict(resp_dict, resp, label, logger)
    return None


if __name__ == '__main__': 
    execute_loc()