# 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

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

# Epoch durations
fixation  = 750
f1        = 500
delay_min = 3000 - 300
delay_max = 3000 + 300
f2        = 500
decision  = 500
tmax      = fixation + f1 + delay_max + f2 + decision

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

    delay = context.get('delay')
    if delay is None:
        delay = tasktools.uniform(rng, dt, delay_min, delay_max)

    durations = {
        'fixation':   (0, fixation),
        'f1':         (fixation, fixation + f1),
        'delay':      (fixation + f1, fixation + f1 + delay),
        'f2':         (fixation + f1 + delay, fixation + f1 + delay + f2),
        'decision':   (fixation + f1 + delay + f2, tmax),
        '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/romo.py
import numpy as np
from pyrl import tasktools


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

# Epoch durations
fixation  = 750
f1        = 500
delay_min = 3000 - 300
delay_max = 3000 + 300
f2        = 500
decision  = 500
tmax      = fixation + f1 + delay_max + f2 + decision

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

    delay = context.get('delay')
    if delay is None:
        delay = tasktools.uniform(rng, dt, delay_min, delay_max)

    durations = {
        'fixation':   (0, fixation),
        'f1':         (fixation, fixation + f1),
        'delay':      (fixation + f1, fixation + f1 + delay),
        'f2':         (fixation + f1 + delay, fixation + f1 + delay + f2),
        'decision':   (fixation + f1 + delay + f2, tmax),
        'tmax':       tmax
        }

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