__author__ = 'nick'
import numpy as np


def generate_stimuli(reward_schedule, n_epochs=20, n_stim_per_epoch=20, n_stim = 2):
    """

    :param reward_schedule: tuple with number of rewarded trials (out of 20) for both stimuli for "correct" responses
    :return:
    """

    # The number of times the "correct" side is rewarded (out of 10). Creates 2
    # stimuli per reward scheduale

    # st = list()
    # rs = list()
    # for s in range(n_stim):
    #     st.append(np.ones(n_stim_per_epoch, dtype=int)*s)
    #     st.append(np.concatenate(
    #         [np.zeros(reward_schedule[0], dtype=int),
    #          np.ones(20-reward_schedule[0], dtype=int)]))

    number_non_rewarded = n_stim_per_epoch - reward_schedule[0] - reward_schedule[1]
    
    st = np.concatenate([np.zeros(n_stim_per_epoch, dtype=int), np.ones(n_stim_per_epoch, dtype=int)])
    if number_non_rewarded == 0:
        rs = np.concatenate([
            np.zeros(reward_schedule[0], dtype=int), np.ones(n_stim_per_epoch-reward_schedule[0], dtype=int),
            np.zeros(reward_schedule[1], dtype=int), np.ones(n_stim_per_epoch-reward_schedule[1], dtype=int)
        ])
    else:
        rs = np.concatenate([
                np.zeros(reward_schedule[0], dtype=int), 
                np.ones(n_stim_per_epoch-reward_schedule[0]-number_non_rewarded, dtype=int), 
                np.ones(number_non_rewarded, dtype=int)*-1,
                np.zeros(reward_schedule[1], dtype=int), 
                np.ones(n_stim_per_epoch-reward_schedule[1]-number_non_rewarded, dtype=int),
                np.ones(number_non_rewarded, dtype=int)*-1,
            ])

    idx = np.random.permutation(len(st))
    rs = rs[idx]
    st = st[idx]

    reversal_length = n_epochs / 2
    trials_per_epoch = n_stim_per_epoch*len(set(st))
    n_trials = n_epochs * trials_per_epoch

    stimuli = list()  # preallocate for speed
    rewards = list()  #

    for epoch in range(reversal_length):
        # idx = np.random.permutation(range(trials_per_epoch))
        #
        # stimuli.extend(st[idx].tolist())
        # rewards.extend(rs[idx].tolist())
        stimuli.extend(st.tolist())
        rewards.extend(rs.tolist())

    #rs = (rs - 1) * -1  # reverse the reward schedule
    rs[rs==0] = 2
    rs[rs==1] = 0
    rs[rs==2] = 1
    for epoch in range(reversal_length):
        # idx = np.random.permutation(range(trials_per_epoch))
        #
        # stimuli.extend(st[idx].tolist())
        # rewards.extend(rs[idx].tolist())
        stimuli.extend(st.tolist())
        rewards.extend(rs.tolist())

    stimuli = np.asarray(stimuli)
    rewards = np.asarray(rewards)

    return stimuli, rewards