# Path: pyrl/tasktools.py
# def generate_ei(N, pE=0.8):
# def to_map(*args):
# def get_idx(t, (start, end)):
# def get_epochs_idx(dt, epochs):
# def choice(rng, a):
# def unravel_index(i, dims):
# def uniform(rng, dt, xmin, xmax):
# def truncated_exponential(rng, dt, mean, xmin=0, xmax=np.inf):
# def divide(x, y):
# def correct_2AFC(perf):
#     def __init__(self):
#     EXC = idx[:Nexc]
#     INH = idx[Nexc:]
# class Task(object):


import numpy as np
from pyrl import tasktools

# Training
n_gradient   = n_conditions
n_validation = 100*n_conditions

# Input noise
sigma = np.sqrt(2*100*0.015)

# Reward discount
tau_reward = 10000

# Durations
fixation = 750
stimulus = 2000
tmax     = fixation + stimulus

# Rewards
R_ABORTED = -1
R_CORRECT = +1

def get_condition(rng, dt, context={}):
    #-------------------------------------------------------------------------------------
    # Epochs
    #-------------------------------------------------------------------------------------

    durations = {
        'fixation': (0, fixation),
        'stimulus': (fixation, fixation + stimulus),
        'tmax':     tmax
        }

# Path: pyrl/tasktools.py
# def generate_ei(N, pE=0.8):
# def to_map(*args):
# def get_idx(t, (start, end)):
# def get_epochs_idx(dt, epochs):
# def choice(rng, a):
# def unravel_index(i, dims):
# def uniform(rng, dt, xmin, xmax):
# def truncated_exponential(rng, dt, mean, xmin=0, xmax=np.inf):
# def divide(x, y):
# def correct_2AFC(perf):
#     def __init__(self):
#     EXC = idx[:Nexc]
#     INH = idx[Nexc:]
# class Task(object):


# Path: examples/models/rdm_rt.py
import numpy as np
from pyrl import tasktools


# Training
n_gradient   = n_conditions
n_validation = 100*n_conditions

# Input noise
sigma = np.sqrt(2*100*0.015)

# Reward discount
tau_reward = 10000

# Durations
fixation = 750
stimulus = 2000
tmax     = fixation + stimulus

# Rewards
R_ABORTED = -1
R_CORRECT = +1

def get_condition(rng, dt, context={}):
    #-------------------------------------------------------------------------------------
    # Epochs
    #-------------------------------------------------------------------------------------

    durations = {
        'fixation': (0, fixation),
        'stimulus': (fixation, fixation + stimulus),
        'tmax':     tmax
        }

    time, epochs = tasktools.get_epochs_idx(dt, durations)