for n in np.arange(N_sim):

    print('SIMULATION ', n + 1)

    model = AuGMEnT(S, R, M, A, alpha, beta, discount, eps, g, leak, rew,
                    dic_stim, dic_resp, prop)

    E[n, :], conv_ep[n], REW[n, :] = model.training_12AX(
        N_trial, p_target, criterion, stop, verb, policy, stoc, t_weighted,
        e_weighted)

    print('\t CONVERGED AT TRIAL ', conv_ep[n])

    if do_test:
        N_test = 1000
        perc_expl[n], perc_no_expl[n], perc_soft[n] = model.test(
            N_test, p_target)
        print('Percentage of correct trials during test (exploration): ',
              perc_expl[n], '%')
        print('Percentage of correct trials during test (no exploration): ',
              perc_no_expl[n], '%')
        print('Percentage of correct trials during test (softmax): ',
              perc_soft[n], '%')

folder = 'DATA'
str_conv = folder + '/' + policy_str + '/conv_long_2.txt'
np.savetxt(str_conv, conv_ep)
E_mean = np.mean(np.reshape(E, (N_sim, -1, 50)), axis=2)
str_err = folder + '/' + policy_str + '/error_long_2.txt'
np.savetxt(str_err, E_mean)
R_mean = np.mean(np.reshape(REW, (N_sim, -1, 50)), axis=2)
str_r = folder + '/' + policy_str + '/reward_long_2.txt'
criterion = 'strong'

do_test = True

for n in np.arange(N_sim):

    print('SIMULATION ', n + 1)

    model = AuGMEnT(S, R, M, A, alpha, beta, discount, eps, g, leak, rew,
                    dic_stim, dic_resp, prop)

    E[n, :], conv_ep[n] = model.training_12AX(N_trial, p_target, criterion,
                                              stop)

    print('\t CONVERGED AT TRIAL ', conv_ep[n])

    if do_test:
        N_test = 1000
        perc[n] = model.test(N_test, p_target)
        print('Percentage of correct trials during test: ', perc, '%')

folder = 'DATA'
str_conv = folder + '/AuGMEnT_' + task + '_conv.txt'
np.savetxt(str_conv, conv_ep)
E_mean = np.mean(np.reshape(E, (-1, 50)), axis=1)
str_err = folder + '/AuGMEnT' + task + '_error.txt'
np.savetxt(str_err, E_mean)
if do_test:
    str_perc = folder + '/AuGMEnT' + task + '_perc.txt'
    np.savetxt(str_perc, perc)