from __future__ import print_function
import os
import csv
import numpy as np
from glob import glob
from ptsa.data.readers import BaseEventReader


def calculate_blink_rate(events, return_percent=False):
    """
    Calculates a participant's blink rate based on an events structure. This requires alignment and artifact detection
    to have already been run on the session's EEG data. The blink rate is defined as the fraction of presentation events
    during which the participant blinked or showed other EOG artifacts while the presented item was on the screen. For
    sessions where some presentation events lack EEG data, only the presentation events with data are counted.

    :param events: An events structure
    :param pres_duration: The number of milliseconds for which each item was presented on the screen (1600 for ltpFR2)
    :param return_percent: If true, returns the blink rate as a percentage. If false, returns the blink rate as a ratio.
    :return: The participant's left-eye, right-eye, and combined blink rates, or np.nan if no presentation events with
    EEG data are available.
    """
    # Get only word presentation events with eeg
    pres_events = events[np.logical_and(events['type'] == 'WORD', np.logical_not(events['eegfile'] == ''))]
    # Count number of word presentations with artifactMS info
    total_pres = float(pres_events.shape[0])
    # If there are no presentation events with artifactMS info, blink rate cannot be calculated
    if total_pres == 0:
        return np.nan

    # Count number of word presentations where eye movements were detected while the word was on the screen
    pres_with_left_blink = np.logical_or(pres_events['eogArtifact'] == 1, pres_events['eogArtifact'] == 3).sum()
    pres_with_right_blink = np.logical_or(pres_events['eogArtifact'] == 2, pres_events['eogArtifact'] == 3).sum()
    pres_with_blink = np.sum(pres_events['eogArtifact'] > 0)

    # Calculate blink rate as presentations with blinks / total presentations
    lbr = pres_with_left_blink / total_pres if total_pres > 0 else np.nan
    rbr = pres_with_right_blink / total_pres if total_pres > 0 else np.nan
    br = pres_with_blink / total_pres if total_pres > 0 else np.nan

    if return_percent:
        return lbr * 100, rbr * 100, br * 100
    else:
        return lbr, rbr, br


def get_math_correct(subj, sess):
    """
    Looks up a participant's math score from a single session by accessing their session log. The total math score for
    each session should be listed near the end of that session's log. If no score can be found in the session log or no
    session log exists, returns np.nan.

    :param subj: The subject ID of the target participant
    :param sess: The session number, as an integer (0-indexed)
    :return: The total number of math questions the participant answered correctly during the specified session, or
    np.nan if the math score cannot be found.
    """
    # Load the session log from that session
    session_log = '/data/eeg/scalp/ltp/ltpFR2/%s/session_%d/session.log' % (subj, sess)
    if os.path.exists(session_log):
        with open(session_log, 'r') as f:
            sess_log = list(csv.reader(f, delimiter='\t'))
        # Find the row that lists the total math score at the end of the session, and extract the math score from it
        total = [r[3] for r in sess_log if len(r) >= 4 and r[2] == 'MATH_TOTAL_SCORE']
        mc = int(total[-1]) if len(total) > 0 else np.nan
    else:
        print('Session log %s cannot be found! Leaving math score as NaN!' % session_log)
        mc = np.nan
    return mc


def calculate_bonus_ltpFR2(subj):
    """
    Calculates bonus payments for each of a participant's 24 sessions based on the following performance brackets:

    P-Recs:
    $0 --> 0% - 19.99%
    $1 --> 20% - 29.99%
    $2 --> 30% - 39.99%
    $3 --> 40% - 49.99%
    $4 --> 50% - 69.99%
    $5 --> 70% - 100%

    Blink rates:
    $0 --> > 50%
    $1 --> 40% - 49.99%
    $2 --> 30% - 39.99%
    $3 --> 20% - 29.99%
    $4 --> 10% - 19.99%
    $5 --> 0% - 9.99%

    Math Correct:
    $0 --> < 200
    $1 --> 200 - 349
    $2 --> 350 - 399
    $3 --> 400 - 449
    $4 --> 450 - 499
    $5 --> > 500

    Recall scores and bonuses can only be calculated once the session has been annotated. Blink rates can only be
    calculated if the session has been successfully aligned and blink detection has been run. If not all presentation
    events have EEG data, the blink rate is calculated only over the events that do.

    :param subj: A string containing the subject ID of the participant for whom to calculate bonuses.
    :return: Returns two numpy arrays. The first is a session x score matrix, with prec in column 0, blink rates in
    columns 1-3, and math score in column 4. The second is a session x bonus matrix, with recall bonus in column 0,
    blink bonus in column 1, math bonus in column 2, and total bonus in column 3.
    """

    # Set experiment parameters and performance bracket boundaries
    n_sessions = 24
    brackets = dict(
        prec=[20, 30, 40, 50, 70],
        br=[10, 20, 30, 40, 50],
        mc=[200, 350, 400, 450, 500]
    )

    scores = np.zeros((24, 5))
    bonuses = np.zeros((24, 4))
    # Calculate scores and bonuses for each session
    for sess in range(n_sessions):
        print(subj, sess)
        # If session has exists and has been post-processed, calculate prec and blink rate, otherwise, set as nan
        event_file = '/protocols/ltp/subjects/%s/experiments/ltpFR2/sessions/%d/behavioral/current_processed/task_events.json' % (subj, sess)
        prec = np.nan
        lbr = np.nan
        rbr = np.nan
        br = np.nan
        try:
            # Calculate performance from the target session
            sess_dir = '/data/eeg/scalp/ltp/ltpFR2/%s/session_%d/' % (subj, sess)
            sess_precs = []
            lsts = glob(os.path.join(sess_dir, '*.lst'))
            for lst in lsts:
                par = os.path.splitext(lst)[0] + '.par'
                with open(lst, 'r') as f:
                    pres = [w.strip() for w in f.readlines()]
                if os.path.exists(par):
                    with open(par, 'r') as f:
                        rec = [w.split('\t')[2].strip() for w in f.readlines()]
                        recalled = [(w in rec) for w in pres]
                        sess_precs.append(np.mean(recalled))
                else:
                    sess_precs.append(np.nan)
            prec = np.nanmean(sess_precs) * 100

            # Load events for given subject and session
            ev = BaseEventReader(filename=event_file, common_root='data', eliminate_nans=False, eliminate_events_with_no_eeg=False).read()
            lbr, rbr, br = calculate_blink_rate(ev, return_percent=True)
            del ev
        except Exception as e:
            # Exceptions here are caused by a nonexistent, empty, or otherwise unreadable event file.
            print(e)
            print('PTSA was unable to read event file %s... Leaving blink rate and recall probability as NaN!' % event_file)

        # Math score can be calculated even before the session has been post-processed, as it only relies on the log
        mc = get_math_correct(subj, sess)

        # Calculate bonuses based on performance brackets
        prec_bonus = np.searchsorted(brackets['prec'], prec, side='right') if not np.isnan(prec) else np.nan
        blink_bonus = 5 - np.searchsorted(brackets['br'], br, side='right') if not np.isnan(br) else np.nan
        math_bonus = np.searchsorted(brackets['mc'], mc, side='right') if not np.isnan(mc) else np.nan
        total_bonus = prec_bonus + blink_bonus + math_bonus

        # Record scores and bonuses from session
        scores[sess] = [prec, round(lbr, 1), round(rbr, 1), round(br, 1), mc]
        bonuses[sess] = [prec_bonus, blink_bonus, math_bonus, total_bonus]

    return scores, bonuses

737f93b2f73d41ae82c97e1b3820a0dba5e4cc12
<mask token>


def calculate_blink_rate(events, return_percent=False):
    """
    Calculates a participant's blink rate based on an events structure. This requires alignment and artifact detection
    to have already been run on the session's EEG data. The blink rate is defined as the fraction of presentation events
    during which the participant blinked or showed other EOG artifacts while the presented item was on the screen. For
    sessions where some presentation events lack EEG data, only the presentation events with data are counted.

    :param events: An events structure
    :param pres_duration: The number of milliseconds for which each item was presented on the screen (1600 for ltpFR2)
    :param return_percent: If true, returns the blink rate as a percentage. If false, returns the blink rate as a ratio.
    :return: The participant's left-eye, right-eye, and combined blink rates, or np.nan if no presentation events with
    EEG data are available.
    """
    pres_events = events[np.logical_and(events['type'] == 'WORD', np.
        logical_not(events['eegfile'] == ''))]
    total_pres = float(pres_events.shape[0])
    if total_pres == 0:
        return np.nan
    pres_with_left_blink = np.logical_or(pres_events['eogArtifact'] == 1, 
        pres_events['eogArtifact'] == 3).sum()
    pres_with_right_blink = np.logical_or(pres_events['eogArtifact'] == 2, 
        pres_events['eogArtifact'] == 3).sum()
    pres_with_blink = np.sum(pres_events['eogArtifact'] > 0)
    lbr = pres_with_left_blink / total_pres if total_pres > 0 else np.nan
    rbr = pres_with_right_blink / total_pres if total_pres > 0 else np.nan
    br = pres_with_blink / total_pres if total_pres > 0 else np.nan
    if return_percent:
        return lbr * 100, rbr * 100, br * 100
    else:
        return lbr, rbr, br


<mask token>

<mask token>


def calculate_blink_rate(events, return_percent=False):
    """
    Calculates a participant's blink rate based on an events structure. This requires alignment and artifact detection
    to have already been run on the session's EEG data. The blink rate is defined as the fraction of presentation events
    during which the participant blinked or showed other EOG artifacts while the presented item was on the screen. For
    sessions where some presentation events lack EEG data, only the presentation events with data are counted.

    :param events: An events structure
    :param pres_duration: The number of milliseconds for which each item was presented on the screen (1600 for ltpFR2)
    :param return_percent: If true, returns the blink rate as a percentage. If false, returns the blink rate as a ratio.
    :return: The participant's left-eye, right-eye, and combined blink rates, or np.nan if no presentation events with
    EEG data are available.
    """
    pres_events = events[np.logical_and(events['type'] == 'WORD', np.
        logical_not(events['eegfile'] == ''))]
    total_pres = float(pres_events.shape[0])
    if total_pres == 0:
        return np.nan
    pres_with_left_blink = np.logical_or(pres_events['eogArtifact'] == 1, 
        pres_events['eogArtifact'] == 3).sum()
    pres_with_right_blink = np.logical_or(pres_events['eogArtifact'] == 2, 
        pres_events['eogArtifact'] == 3).sum()
    pres_with_blink = np.sum(pres_events['eogArtifact'] > 0)
    lbr = pres_with_left_blink / total_pres if total_pres > 0 else np.nan
    rbr = pres_with_right_blink / total_pres if total_pres > 0 else np.nan
    br = pres_with_blink / total_pres if total_pres > 0 else np.nan
    if return_percent:
        return lbr * 100, rbr * 100, br * 100
    else:
        return lbr, rbr, br


<mask token>


def calculate_bonus_ltpFR2(subj):
    """
    Calculates bonus payments for each of a participant's 24 sessions based on the following performance brackets:

    P-Recs:
    $0 --> 0% - 19.99%
    $1 --> 20% - 29.99%
    $2 --> 30% - 39.99%
    $3 --> 40% - 49.99%
    $4 --> 50% - 69.99%
    $5 --> 70% - 100%

    Blink rates:
    $0 --> > 50%
    $1 --> 40% - 49.99%
    $2 --> 30% - 39.99%
    $3 --> 20% - 29.99%
    $4 --> 10% - 19.99%
    $5 --> 0% - 9.99%

    Math Correct:
    $0 --> < 200
    $1 --> 200 - 349
    $2 --> 350 - 399
    $3 --> 400 - 449
    $4 --> 450 - 499
    $5 --> > 500

    Recall scores and bonuses can only be calculated once the session has been annotated. Blink rates can only be
    calculated if the session has been successfully aligned and blink detection has been run. If not all presentation
    events have EEG data, the blink rate is calculated only over the events that do.

    :param subj: A string containing the subject ID of the participant for whom to calculate bonuses.
    :return: Returns two numpy arrays. The first is a session x score matrix, with prec in column 0, blink rates in
    columns 1-3, and math score in column 4. The second is a session x bonus matrix, with recall bonus in column 0,
    blink bonus in column 1, math bonus in column 2, and total bonus in column 3.
    """
    n_sessions = 24
    brackets = dict(prec=[20, 30, 40, 50, 70], br=[10, 20, 30, 40, 50], mc=
        [200, 350, 400, 450, 500])
    scores = np.zeros((24, 5))
    bonuses = np.zeros((24, 4))
    for sess in range(n_sessions):
        print(subj, sess)
        event_file = (
            '/protocols/ltp/subjects/%s/experiments/ltpFR2/sessions/%d/behavioral/current_processed/task_events.json'
             % (subj, sess))
        prec = np.nan
        lbr = np.nan
        rbr = np.nan
        br = np.nan
        try:
            sess_dir = '/data/eeg/scalp/ltp/ltpFR2/%s/session_%d/' % (subj,
                sess)
            sess_precs = []
            lsts = glob(os.path.join(sess_dir, '*.lst'))
            for lst in lsts:
                par = os.path.splitext(lst)[0] + '.par'
                with open(lst, 'r') as f:
                    pres = [w.strip() for w in f.readlines()]
                if os.path.exists(par):
                    with open(par, 'r') as f:
                        rec = [w.split('\t')[2].strip() for w in f.readlines()]
                        recalled = [(w in rec) for w in pres]
                        sess_precs.append(np.mean(recalled))
                else:
                    sess_precs.append(np.nan)
            prec = np.nanmean(sess_precs) * 100
            ev = BaseEventReader(filename=event_file, common_root='data',
                eliminate_nans=False, eliminate_events_with_no_eeg=False).read(
                )
            lbr, rbr, br = calculate_blink_rate(ev, return_percent=True)
            del ev
        except Exception as e:
            print(e)
            print(
                'PTSA was unable to read event file %s... Leaving blink rate and recall probability as NaN!'
                 % event_file)
        mc = get_math_correct(subj, sess)
        prec_bonus = np.searchsorted(brackets['prec'], prec, side='right'
            ) if not np.isnan(prec) else np.nan
        blink_bonus = 5 - np.searchsorted(brackets['br'], br, side='right'
            ) if not np.isnan(br) else np.nan
        math_bonus = np.searchsorted(brackets['mc'], mc, side='right'
            ) if not np.isnan(mc) else np.nan
        total_bonus = prec_bonus + blink_bonus + math_bonus
        scores[sess] = [prec, round(lbr, 1), round(rbr, 1), round(br, 1), mc]
        bonuses[sess] = [prec_bonus, blink_bonus, math_bonus, total_bonus]
    return scores, bonuses

<mask token>


def calculate_blink_rate(events, return_percent=False):
    """
    Calculates a participant's blink rate based on an events structure. This requires alignment and artifact detection
    to have already been run on the session's EEG data. The blink rate is defined as the fraction of presentation events
    during which the participant blinked or showed other EOG artifacts while the presented item was on the screen. For
    sessions where some presentation events lack EEG data, only the presentation events with data are counted.

    :param events: An events structure
    :param pres_duration: The number of milliseconds for which each item was presented on the screen (1600 for ltpFR2)
    :param return_percent: If true, returns the blink rate as a percentage. If false, returns the blink rate as a ratio.
    :return: The participant's left-eye, right-eye, and combined blink rates, or np.nan if no presentation events with
    EEG data are available.
    """
    pres_events = events[np.logical_and(events['type'] == 'WORD', np.
        logical_not(events['eegfile'] == ''))]
    total_pres = float(pres_events.shape[0])
    if total_pres == 0:
        return np.nan
    pres_with_left_blink = np.logical_or(pres_events['eogArtifact'] == 1, 
        pres_events['eogArtifact'] == 3).sum()
    pres_with_right_blink = np.logical_or(pres_events['eogArtifact'] == 2, 
        pres_events['eogArtifact'] == 3).sum()
    pres_with_blink = np.sum(pres_events['eogArtifact'] > 0)
    lbr = pres_with_left_blink / total_pres if total_pres > 0 else np.nan
    rbr = pres_with_right_blink / total_pres if total_pres > 0 else np.nan
    br = pres_with_blink / total_pres if total_pres > 0 else np.nan
    if return_percent:
        return lbr * 100, rbr * 100, br * 100
    else:
        return lbr, rbr, br


def get_math_correct(subj, sess):
    """
    Looks up a participant's math score from a single session by accessing their session log. The total math score for
    each session should be listed near the end of that session's log. If no score can be found in the session log or no
    session log exists, returns np.nan.

    :param subj: The subject ID of the target participant
    :param sess: The session number, as an integer (0-indexed)
    :return: The total number of math questions the participant answered correctly during the specified session, or
    np.nan if the math score cannot be found.
    """
    session_log = '/data/eeg/scalp/ltp/ltpFR2/%s/session_%d/session.log' % (
        subj, sess)
    if os.path.exists(session_log):
        with open(session_log, 'r') as f:
            sess_log = list(csv.reader(f, delimiter='\t'))
        total = [r[3] for r in sess_log if len(r) >= 4 and r[2] ==
            'MATH_TOTAL_SCORE']
        mc = int(total[-1]) if len(total) > 0 else np.nan
    else:
        print('Session log %s cannot be found! Leaving math score as NaN!' %
            session_log)
        mc = np.nan
    return mc


def calculate_bonus_ltpFR2(subj):
    """
    Calculates bonus payments for each of a participant's 24 sessions based on the following performance brackets:

    P-Recs:
    $0 --> 0% - 19.99%
    $1 --> 20% - 29.99%
    $2 --> 30% - 39.99%
    $3 --> 40% - 49.99%
    $4 --> 50% - 69.99%
    $5 --> 70% - 100%

    Blink rates:
    $0 --> > 50%
    $1 --> 40% - 49.99%
    $2 --> 30% - 39.99%
    $3 --> 20% - 29.99%
    $4 --> 10% - 19.99%
    $5 --> 0% - 9.99%

    Math Correct:
    $0 --> < 200
    $1 --> 200 - 349
    $2 --> 350 - 399
    $3 --> 400 - 449
    $4 --> 450 - 499
    $5 --> > 500

    Recall scores and bonuses can only be calculated once the session has been annotated. Blink rates can only be
    calculated if the session has been successfully aligned and blink detection has been run. If not all presentation
    events have EEG data, the blink rate is calculated only over the events that do.

    :param subj: A string containing the subject ID of the participant for whom to calculate bonuses.
    :return: Returns two numpy arrays. The first is a session x score matrix, with prec in column 0, blink rates in
    columns 1-3, and math score in column 4. The second is a session x bonus matrix, with recall bonus in column 0,
    blink bonus in column 1, math bonus in column 2, and total bonus in column 3.
    """
    n_sessions = 24
    brackets = dict(prec=[20, 30, 40, 50, 70], br=[10, 20, 30, 40, 50], mc=
        [200, 350, 400, 450, 500])
    scores = np.zeros((24, 5))
    bonuses = np.zeros((24, 4))
    for sess in range(n_sessions):
        print(subj, sess)
        event_file = (
            '/protocols/ltp/subjects/%s/experiments/ltpFR2/sessions/%d/behavioral/current_processed/task_events.json'
             % (subj, sess))
        prec = np.nan
        lbr = np.nan
        rbr = np.nan
        br = np.nan
        try:
            sess_dir = '/data/eeg/scalp/ltp/ltpFR2/%s/session_%d/' % (subj,
                sess)
            sess_precs = []
            lsts = glob(os.path.join(sess_dir, '*.lst'))
            for lst in lsts:
                par = os.path.splitext(lst)[0] + '.par'
                with open(lst, 'r') as f:
                    pres = [w.strip() for w in f.readlines()]
                if os.path.exists(par):
                    with open(par, 'r') as f:
                        rec = [w.split('\t')[2].strip() for w in f.readlines()]
                        recalled = [(w in rec) for w in pres]
                        sess_precs.append(np.mean(recalled))
                else:
                    sess_precs.append(np.nan)
            prec = np.nanmean(sess_precs) * 100
            ev = BaseEventReader(filename=event_file, common_root='data',
                eliminate_nans=False, eliminate_events_with_no_eeg=False).read(
                )
            lbr, rbr, br = calculate_blink_rate(ev, return_percent=True)
            del ev
        except Exception as e:
            print(e)
            print(
                'PTSA was unable to read event file %s... Leaving blink rate and recall probability as NaN!'
                 % event_file)
        mc = get_math_correct(subj, sess)
        prec_bonus = np.searchsorted(brackets['prec'], prec, side='right'
            ) if not np.isnan(prec) else np.nan
        blink_bonus = 5 - np.searchsorted(brackets['br'], br, side='right'
            ) if not np.isnan(br) else np.nan
        math_bonus = np.searchsorted(brackets['mc'], mc, side='right'
            ) if not np.isnan(mc) else np.nan
        total_bonus = prec_bonus + blink_bonus + math_bonus
        scores[sess] = [prec, round(lbr, 1), round(rbr, 1), round(br, 1), mc]
        bonuses[sess] = [prec_bonus, blink_bonus, math_bonus, total_bonus]
    return scores, bonuses

from __future__ import print_function
import os
import csv
import numpy as np
from glob import glob
from ptsa.data.readers import BaseEventReader


def calculate_blink_rate(events, return_percent=False):
    """
    Calculates a participant's blink rate based on an events structure. This requires alignment and artifact detection
    to have already been run on the session's EEG data. The blink rate is defined as the fraction of presentation events
    during which the participant blinked or showed other EOG artifacts while the presented item was on the screen. For
    sessions where some presentation events lack EEG data, only the presentation events with data are counted.

    :param events: An events structure
    :param pres_duration: The number of milliseconds for which each item was presented on the screen (1600 for ltpFR2)
    :param return_percent: If true, returns the blink rate as a percentage. If false, returns the blink rate as a ratio.
    :return: The participant's left-eye, right-eye, and combined blink rates, or np.nan if no presentation events with
    EEG data are available.
    """
    pres_events = events[np.logical_and(events['type'] == 'WORD', np.
        logical_not(events['eegfile'] == ''))]
    total_pres = float(pres_events.shape[0])
    if total_pres == 0:
        return np.nan
    pres_with_left_blink = np.logical_or(pres_events['eogArtifact'] == 1, 
        pres_events['eogArtifact'] == 3).sum()
    pres_with_right_blink = np.logical_or(pres_events['eogArtifact'] == 2, 
        pres_events['eogArtifact'] == 3).sum()
    pres_with_blink = np.sum(pres_events['eogArtifact'] > 0)
    lbr = pres_with_left_blink / total_pres if total_pres > 0 else np.nan
    rbr = pres_with_right_blink / total_pres if total_pres > 0 else np.nan
    br = pres_with_blink / total_pres if total_pres > 0 else np.nan
    if return_percent:
        return lbr * 100, rbr * 100, br * 100
    else:
        return lbr, rbr, br


def get_math_correct(subj, sess):
    """
    Looks up a participant's math score from a single session by accessing their session log. The total math score for
    each session should be listed near the end of that session's log. If no score can be found in the session log or no
    session log exists, returns np.nan.

    :param subj: The subject ID of the target participant
    :param sess: The session number, as an integer (0-indexed)
    :return: The total number of math questions the participant answered correctly during the specified session, or
    np.nan if the math score cannot be found.
    """
    session_log = '/data/eeg/scalp/ltp/ltpFR2/%s/session_%d/session.log' % (
        subj, sess)
    if os.path.exists(session_log):
        with open(session_log, 'r') as f:
            sess_log = list(csv.reader(f, delimiter='\t'))
        total = [r[3] for r in sess_log if len(r) >= 4 and r[2] ==
            'MATH_TOTAL_SCORE']
        mc = int(total[-1]) if len(total) > 0 else np.nan
    else:
        print('Session log %s cannot be found! Leaving math score as NaN!' %
            session_log)
        mc = np.nan
    return mc


def calculate_bonus_ltpFR2(subj):
    """
    Calculates bonus payments for each of a participant's 24 sessions based on the following performance brackets:

    P-Recs:
    $0 --> 0% - 19.99%
    $1 --> 20% - 29.99%
    $2 --> 30% - 39.99%
    $3 --> 40% - 49.99%
    $4 --> 50% - 69.99%
    $5 --> 70% - 100%

    Blink rates:
    $0 --> > 50%
    $1 --> 40% - 49.99%
    $2 --> 30% - 39.99%
    $3 --> 20% - 29.99%
    $4 --> 10% - 19.99%
    $5 --> 0% - 9.99%

    Math Correct:
    $0 --> < 200
    $1 --> 200 - 349
    $2 --> 350 - 399
    $3 --> 400 - 449
    $4 --> 450 - 499
    $5 --> > 500

    Recall scores and bonuses can only be calculated once the session has been annotated. Blink rates can only be
    calculated if the session has been successfully aligned and blink detection has been run. If not all presentation
    events have EEG data, the blink rate is calculated only over the events that do.

    :param subj: A string containing the subject ID of the participant for whom to calculate bonuses.
    :return: Returns two numpy arrays. The first is a session x score matrix, with prec in column 0, blink rates in
    columns 1-3, and math score in column 4. The second is a session x bonus matrix, with recall bonus in column 0,
    blink bonus in column 1, math bonus in column 2, and total bonus in column 3.
    """
    n_sessions = 24
    brackets = dict(prec=[20, 30, 40, 50, 70], br=[10, 20, 30, 40, 50], mc=
        [200, 350, 400, 450, 500])
    scores = np.zeros((24, 5))
    bonuses = np.zeros((24, 4))
    for sess in range(n_sessions):
        print(subj, sess)
        event_file = (
            '/protocols/ltp/subjects/%s/experiments/ltpFR2/sessions/%d/behavioral/current_processed/task_events.json'
             % (subj, sess))
        prec = np.nan
        lbr = np.nan
        rbr = np.nan
        br = np.nan
        try:
            sess_dir = '/data/eeg/scalp/ltp/ltpFR2/%s/session_%d/' % (subj,
                sess)
            sess_precs = []
            lsts = glob(os.path.join(sess_dir, '*.lst'))
            for lst in lsts:
                par = os.path.splitext(lst)[0] + '.par'
                with open(lst, 'r') as f:
                    pres = [w.strip() for w in f.readlines()]
                if os.path.exists(par):
                    with open(par, 'r') as f:
                        rec = [w.split('\t')[2].strip() for w in f.readlines()]
                        recalled = [(w in rec) for w in pres]
                        sess_precs.append(np.mean(recalled))
                else:
                    sess_precs.append(np.nan)
            prec = np.nanmean(sess_precs) * 100
            ev = BaseEventReader(filename=event_file, common_root='data',
                eliminate_nans=False, eliminate_events_with_no_eeg=False).read(
                )
            lbr, rbr, br = calculate_blink_rate(ev, return_percent=True)
            del ev
        except Exception as e:
            print(e)
            print(
                'PTSA was unable to read event file %s... Leaving blink rate and recall probability as NaN!'
                 % event_file)
        mc = get_math_correct(subj, sess)
        prec_bonus = np.searchsorted(brackets['prec'], prec, side='right'
            ) if not np.isnan(prec) else np.nan
        blink_bonus = 5 - np.searchsorted(brackets['br'], br, side='right'
            ) if not np.isnan(br) else np.nan
        math_bonus = np.searchsorted(brackets['mc'], mc, side='right'
            ) if not np.isnan(mc) else np.nan
        total_bonus = prec_bonus + blink_bonus + math_bonus
        scores[sess] = [prec, round(lbr, 1), round(rbr, 1), round(br, 1), mc]
        bonuses[sess] = [prec_bonus, blink_bonus, math_bonus, total_bonus]
    return scores, bonuses

from __future__ import print_function
import os
import csv
import numpy as np
from glob import glob
from ptsa.data.readers import BaseEventReader


def calculate_blink_rate(events, return_percent=False):
    """
    Calculates a participant's blink rate based on an events structure. This requires alignment and artifact detection
    to have already been run on the session's EEG data. The blink rate is defined as the fraction of presentation events
    during which the participant blinked or showed other EOG artifacts while the presented item was on the screen. For
    sessions where some presentation events lack EEG data, only the presentation events with data are counted.

    :param events: An events structure
    :param pres_duration: The number of milliseconds for which each item was presented on the screen (1600 for ltpFR2)
    :param return_percent: If true, returns the blink rate as a percentage. If false, returns the blink rate as a ratio.
    :return: The participant's left-eye, right-eye, and combined blink rates, or np.nan if no presentation events with
    EEG data are available.
    """
    # Get only word presentation events with eeg
    pres_events = events[np.logical_and(events['type'] == 'WORD', np.logical_not(events['eegfile'] == ''))]
    # Count number of word presentations with artifactMS info
    total_pres = float(pres_events.shape[0])
    # If there are no presentation events with artifactMS info, blink rate cannot be calculated
    if total_pres == 0:
        return np.nan

    # Count number of word presentations where eye movements were detected while the word was on the screen
    pres_with_left_blink = np.logical_or(pres_events['eogArtifact'] == 1, pres_events['eogArtifact'] == 3).sum()
    pres_with_right_blink = np.logical_or(pres_events['eogArtifact'] == 2, pres_events['eogArtifact'] == 3).sum()
    pres_with_blink = np.sum(pres_events['eogArtifact'] > 0)

    # Calculate blink rate as presentations with blinks / total presentations
    lbr = pres_with_left_blink / total_pres if total_pres > 0 else np.nan
    rbr = pres_with_right_blink / total_pres if total_pres > 0 else np.nan
    br = pres_with_blink / total_pres if total_pres > 0 else np.nan

    if return_percent:
        return lbr * 100, rbr * 100, br * 100
    else:
        return lbr, rbr, br


def get_math_correct(subj, sess):
    """
    Looks up a participant's math score from a single session by accessing their session log. The total math score for
    each session should be listed near the end of that session's log. If no score can be found in the session log or no
    session log exists, returns np.nan.

    :param subj: The subject ID of the target participant
    :param sess: The session number, as an integer (0-indexed)
    :return: The total number of math questions the participant answered correctly during the specified session, or
    np.nan if the math score cannot be found.
    """
    # Load the session log from that session
    session_log = '/data/eeg/scalp/ltp/ltpFR2/%s/session_%d/session.log' % (subj, sess)
    if os.path.exists(session_log):
        with open(session_log, 'r') as f:
            sess_log = list(csv.reader(f, delimiter='\t'))
        # Find the row that lists the total math score at the end of the session, and extract the math score from it
        total = [r[3] for r in sess_log if len(r) >= 4 and r[2] == 'MATH_TOTAL_SCORE']
        mc = int(total[-1]) if len(total) > 0 else np.nan
    else:
        print('Session log %s cannot be found! Leaving math score as NaN!' % session_log)
        mc = np.nan
    return mc


def calculate_bonus_ltpFR2(subj):
    """
    Calculates bonus payments for each of a participant's 24 sessions based on the following performance brackets:

    P-Recs:
    $0 --> 0% - 19.99%
    $1 --> 20% - 29.99%
    $2 --> 30% - 39.99%
    $3 --> 40% - 49.99%
    $4 --> 50% - 69.99%
    $5 --> 70% - 100%

    Blink rates:
    $0 --> > 50%
    $1 --> 40% - 49.99%
    $2 --> 30% - 39.99%
    $3 --> 20% - 29.99%
    $4 --> 10% - 19.99%
    $5 --> 0% - 9.99%

    Math Correct:
    $0 --> < 200
    $1 --> 200 - 349
    $2 --> 350 - 399
    $3 --> 400 - 449
    $4 --> 450 - 499
    $5 --> > 500

    Recall scores and bonuses can only be calculated once the session has been annotated. Blink rates can only be
    calculated if the session has been successfully aligned and blink detection has been run. If not all presentation
    events have EEG data, the blink rate is calculated only over the events that do.

    :param subj: A string containing the subject ID of the participant for whom to calculate bonuses.
    :return: Returns two numpy arrays. The first is a session x score matrix, with prec in column 0, blink rates in
    columns 1-3, and math score in column 4. The second is a session x bonus matrix, with recall bonus in column 0,
    blink bonus in column 1, math bonus in column 2, and total bonus in column 3.
    """

    # Set experiment parameters and performance bracket boundaries
    n_sessions = 24
    brackets = dict(
        prec=[20, 30, 40, 50, 70],
        br=[10, 20, 30, 40, 50],
        mc=[200, 350, 400, 450, 500]
    )

    scores = np.zeros((24, 5))
    bonuses = np.zeros((24, 4))
    # Calculate scores and bonuses for each session
    for sess in range(n_sessions):
        print(subj, sess)
        # If session has exists and has been post-processed, calculate prec and blink rate, otherwise, set as nan
        event_file = '/protocols/ltp/subjects/%s/experiments/ltpFR2/sessions/%d/behavioral/current_processed/task_events.json' % (subj, sess)
        prec = np.nan
        lbr = np.nan
        rbr = np.nan
        br = np.nan
        try:
            # Calculate performance from the target session
            sess_dir = '/data/eeg/scalp/ltp/ltpFR2/%s/session_%d/' % (subj, sess)
            sess_precs = []
            lsts = glob(os.path.join(sess_dir, '*.lst'))
            for lst in lsts:
                par = os.path.splitext(lst)[0] + '.par'
                with open(lst, 'r') as f:
                    pres = [w.strip() for w in f.readlines()]
                if os.path.exists(par):
                    with open(par, 'r') as f:
                        rec = [w.split('\t')[2].strip() for w in f.readlines()]
                        recalled = [(w in rec) for w in pres]
                        sess_precs.append(np.mean(recalled))
                else:
                    sess_precs.append(np.nan)
            prec = np.nanmean(sess_precs) * 100

            # Load events for given subject and session
            ev = BaseEventReader(filename=event_file, common_root='data', eliminate_nans=False, eliminate_events_with_no_eeg=False).read()
            lbr, rbr, br = calculate_blink_rate(ev, return_percent=True)
            del ev
        except Exception as e:
            # Exceptions here are caused by a nonexistent, empty, or otherwise unreadable event file.
            print(e)
            print('PTSA was unable to read event file %s... Leaving blink rate and recall probability as NaN!' % event_file)

        # Math score can be calculated even before the session has been post-processed, as it only relies on the log
        mc = get_math_correct(subj, sess)

        # Calculate bonuses based on performance brackets
        prec_bonus = np.searchsorted(brackets['prec'], prec, side='right') if not np.isnan(prec) else np.nan
        blink_bonus = 5 - np.searchsorted(brackets['br'], br, side='right') if not np.isnan(br) else np.nan
        math_bonus = np.searchsorted(brackets['mc'], mc, side='right') if not np.isnan(mc) else np.nan
        total_bonus = prec_bonus + blink_bonus + math_bonus

        # Record scores and bonuses from session
        scores[sess] = [prec, round(lbr, 1), round(rbr, 1), round(br, 1), mc]
        bonuses[sess] = [prec_bonus, blink_bonus, math_bonus, total_bonus]

    return scores, bonuses
