"""
Perceptual decision-making with postdecision wagering, based on

  Representation of confidence associated with a decision by
  neurons in the parietal cortex.
  R. Kiani & M. N. Shadlen, Science 2009.

  http://dx.doi.org/10.1126/science.1169405

"""
from __future__ import division

import numpy as np

from pyrl import tasktools

# Inputs
inputs = tasktools.to_map('FIXATION', 'LEFT', 'RIGHT', 'SURE')

# Actions
actions = tasktools.to_map('FIXATE', 'CHOOSE-LEFT', 'CHOOSE-RIGHT', 'CHOOSE-SURE')

# Trial conditions
wagers       = [True, False]
left_rights  = [-1, 1]
cohs         = [0, 3.2, 6.4, 12.8, 25.6, 51.2]
n_conditions = len(wagers) * len(left_rights) * len(cohs)

# Training
n_gradient   = n_conditions
n_validation = 50*n_conditions

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

N = 200

# Separate inputs
N    = 100
Wins = []
for i in xrange(3):
    Win = np.zeros((len(inputs), N))
    Win[inputs['FIXATION']]      = 1
    Win[inputs['LEFT'], :N//2] = 1
    Win[inputs['RIGHT'],:N//2] = 1
    Win[inputs['SURE'], N//2:] = 1
    Wins.append(Win)
Win = np.concatenate(Wins, axis=1)

# Durations
fixation      = 750
stimulus_min  = 100
stimulus_mean = 180
stimulus_max  = 800
delay_min     = 1200
delay_mean    = 1350
delay_max     = 1800
sure_min      = 500
sure_mean     = 575
sure_max      = 750
decision      = 500
tmax          = fixation + stimulus_min + stimulus_max + delay_max + decision

# 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,
                                                xmin=delay_min, xmax=delay_max)

    if wager:
        sure_onset = context.get('sure_onset')
        if sure_onset is None:
            sure_onset = tasktools.truncated_exponential(rng, dt, sure_mean,
                                                         xmin=sure_min, xmax=sure_max)

    durations = {
        'fixation':  (0, fixation),
        'stimulus':  (fixation, fixation + stimulus),
        'delay':     (fixation + stimulus, fixation + stimulus + delay),
        'decision':  (fixation + stimulus + delay, tmax),
        'tmax':      tmax
        }
    if wager:
        durations['sure'] = (fixation + stimulus + sure_onset, tmax)
    time, epochs = tasktools.get_epochs_idx(dt, durations)

    #-------------------------------------------------------------------------------------
    # Trial
    #-------------------------------------------------------------------------------------

    left_right = context.get('left_right')
    if left_right is None:
        left_right = rng.choice(left_rights)

    coh = context.get('coh')
    if coh is None:
        coh = rng.choice(cohs)

    return {
        'durations':  durations,
        'time':       time,
        'epochs':     epochs,
        'wager':      wager,
        'left_right': left_right,
        'coh':        coh
        }

def get_step(rng, dt, trial, t, a):
    #-------------------------------------------------------------------------------------
    # Reward
    #-------------------------------------------------------------------------------------

    epochs = trial['epochs']
    status = {'continue': True}
    reward = 0
    if t-1 not in epochs['decision']:
        if a != actions['FIXATE']:
            status['continue'] = False
            reward = R_ABORTED
    elif t-1 in epochs['decision']:
        if a == actions['CHOOSE-LEFT']: