# 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: pyrl/performance.py
# class PerformancePostdecisionWager(object):
#     def __init__(self):
#         self.wagers    = []
#         self.corrects  = []
#         self.choices   = []
#         self.t_choices = []
# 
#     def update(self, trial, status):
#         self.wagers.append(trial['wager'])
#         self.corrects.append(status.get('correct'))
#         self.choices.append(status.get('choice'))
#         self.t_choices.append(status.get('t_choice'))
# 
#     @property
#     def n_correct(self):
#         return sum([c for c in self.corrects if c is not None])
# 
#     @property
#     def n_sure_decision(self):
#         return len([1 for w, c in zip(self.wagers, self.choices) if w and c is not None])
# 
#     @property
#     def n_trials(self):
#         return len(self.choices)
# 
#     @property
#     def n_decision(self):
#         return len([1 for c in self.choices if c in ['L', 'R']])
# 
#     @property
#     def n_sure(self):
#         return len([1 for c in self.choices if c == 'S'])
# 
#     @property
#     def n_answer(self):
#         return len([1 for c in self.choices if c is not None])
# 
#     @property
#     def n_wager(self):
#         return sum(self.wagers)
# 
#     def display(self, output=True):
#         n_trials        = self.n_trials
#         n_decision      = self.n_decision
#         n_correct       = self.n_correct
#         n_sure_decision = self.n_sure_decision
#         n_sure          = self.n_sure
#         n_answer        = self.n_answer
#         n_wager         = self.n_wager
# 
#         items = OrderedDict()
#         items['P(answer)'] = '{}/{} = {:.3f}'.format(n_answer, n_trials,
#                                                      n_answer/n_trials)
#         items['P(decision)'] = '{}/{} = {:.3f}'.format(n_decision, n_trials,
#                                                        n_decision/n_trials)
#         if n_decision > 0:
#             items['P(correct|decision)'] = '{}/{} = {:.3f}'.format(n_correct, n_decision,
#                                                                    n_correct/n_decision)
#         items['P(wager trials)'] = '{}/{} = {:.3f}'.format(n_wager, n_trials,
#                                                            n_wager/n_trials)
#         if n_sure_decision > 0:
#             items['P(sure)'] = '{}/{} = {:.3f}'.format(n_sure, n_sure_decision,
#                                                        n_sure/n_sure_decision)
# 
#         if output:
#             utils.print_dict(items)
#         return items


import numpy as np
from pyrl import tasktools
from pyrl.performance import PerformancePostdecisionWager as Performance

# Rewards
R_ABORTED = -1
R_CORRECT = +1
R_SURE    = 0.7*R_CORRECT

# Input scaling
def scale(coh):
    return (1 + coh/100)/2

def get_condition(rng, dt, context={}):
    #-------------------------------------------------------------------------------------
    # Wager or no wager?
    #-------------------------------------------------------------------------------------

    wager = context.get('wager')
    if wager is None:
        wager = rng.choice(wagers)

    #-------------------------------------------------------------------------------------
    # Epochs
    #-------------------------------------------------------------------------------------

    stimulus = context.get('stimulus')
    if stimulus is None:
        stimulus = stimulus_min + tasktools.truncated_exponential(rng, dt, stimulus_mean,
                                                                  xmax=stimulus_max)

    delay = context.get('delay')
    if delay is None:

# 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: pyrl/performance.py
# class PerformancePostdecisionWager(object):
#     def __init__(self):
#         self.wagers    = []
#         self.corrects  = []
#         self.choices   = []
#         self.t_choices = []
# 
#     def update(self, trial, status):
#         self.wagers.append(trial['wager'])
#         self.corrects.append(status.get('correct'))
#         self.choices.append(status.get('choice'))
#         self.t_choices.append(status.get('t_choice'))
# 
#     @property
#     def n_correct(self):
#         return sum([c for c in self.corrects if c is not None])
# 
#     @property
#     def n_sure_decision(self):
#         return len([1 for w, c in zip(self.wagers, self.choices) if w and c is not None])
# 
#     @property
#     def n_trials(self):
#         return len(self.choices)
# 
#     @property
#     def n_decision(self):
#         return len([1 for c in self.choices if c in ['L', 'R']])
# 
#     @property
#     def n_sure(self):
#         return len([1 for c in self.choices if c == 'S'])
# 
#     @property
#     def n_answer(self):
#         return len([1 for c in self.choices if c is not None])
# 
#     @property
#     def n_wager(self):
#         return sum(self.wagers)
# 
#     def display(self, output=True):
#         n_trials        = self.n_trials
#         n_decision      = self.n_decision
#         n_correct       = self.n_correct
#         n_sure_decision = self.n_sure_decision
#         n_sure          = self.n_sure
#         n_answer        = self.n_answer
#         n_wager         = self.n_wager
# 
#         items = OrderedDict()
#         items['P(answer)'] = '{}/{} = {:.3f}'.format(n_answer, n_trials,
#                                                      n_answer/n_trials)
#         items['P(decision)'] = '{}/{} = {:.3f}'.format(n_decision, n_trials,
#                                                        n_decision/n_trials)
#         if n_decision > 0:
#             items['P(correct|decision)'] = '{}/{} = {:.3f}'.format(n_correct, n_decision,
#                                                                    n_correct/n_decision)
#         items['P(wager trials)'] = '{}/{} = {:.3f}'.format(n_wager, n_trials,
#                                                            n_wager/n_trials)
#         if n_sure_decision > 0:
#             items['P(sure)'] = '{}/{} = {:.3f}'.format(n_sure, n_sure_decision,
#                                                        n_sure/n_sure_decision)
# 
#         if output:
#             utils.print_dict(items)
#         return items


# Path: examples/models/postdecisionwager_large.py
import numpy as np
from pyrl import tasktools
from pyrl.performance import PerformancePostdecisionWager as Performance


# Rewards
R_ABORTED = -1
R_CORRECT = +1
R_SURE    = 0.7*R_CORRECT

# Input scaling
def scale(coh):
    return (1 + coh/100)/2

def get_condition(rng, dt, context={}):
    #-------------------------------------------------------------------------------------
    # Wager or no wager?
    #-------------------------------------------------------------------------------------

    wager = context.get('wager')
    if wager is None:
        wager = rng.choice(wagers)

    #-------------------------------------------------------------------------------------
    # Epochs
    #-------------------------------------------------------------------------------------

    stimulus = context.get('stimulus')
    if stimulus is None:
        stimulus = stimulus_min + tasktools.truncated_exponential(rng, dt, stimulus_mean,
                                                                  xmax=stimulus_max)

    delay = context.get('delay')
    if delay is None:

        delay = tasktools.truncated_exponential(rng, dt, delay_mean,