import os
import numpy as np
import scipy as sp
import scipy.stats as stats
import scipy.signal as signal
import matplotlib.pyplot as plt
from sympy import *
import math

%matplotlib inline 

# define an Impulse Response Function:
def pupil_IRF(timepoints, s=1.0/(10**26), n=10.1, tmax=0.93):
    
    """ pupil_IRF defines the IRF (response to a transient input) of the pupil.
    
    Parameters
    ----------
    t: IRF is defined with respect to 't'
    s: scaling factor
    n: sets the width
    tmax: sets the time to peak 
    IRF_len: function is evaluated for t = [0:IRF_len]
    
    Returns
    -------
    y: IRF evaluated for t = [0:IRF_len]
    yprime: IRF first derivative evaluated for t = [0:IRF_len]
    
    """
    
    # in sympy:
    t = Symbol('t')
    y = ( (s) * (t**n) * (math.e**((-n*t)/tmax)) )
    yprime = y.diff(t)
    
    # lambdify:
    y = lambdify(t, y, "numpy")
    yprime = lambdify(t, yprime, "numpy")
    
    # evaluate:
    y = y(timepoints)
    yprime = yprime(timepoints)
    
    return (y, yprime)

# create the IRF:
sample_rate = 10
IRF_len = 3.0 # in seconds
timepoints = np.linspace(0,IRF_len,IRF_len*sample_rate)

IRF, IRF_prime = pupil_IRF(timepoints=timepoints)
IRF = IRF / IRF.std()
IRF_prime = IRF_prime / IRF_prime.std()

# plot the IRF:
fig = plt.figure()
plt.plot(timepoints, IRF, color='r')
# plt.plot(timepoints, IRF_prime, color='g')
plt.legend(['IRF'])
plt.title('Impulse Response Function')
plt.xlabel('time (s)')
plt.ylabel('a.u.')

duration = 60 # in seconds
times_l = np.array([5,15,16,25])
times_r = np.array([35,45,46,55])
input_signal_l = np.zeros(duration * sample_rate)
input_signal_r = np.zeros(duration * sample_rate)
for i in times_l:
    input_signal_l[i*sample_rate] = 1.5
for i in times_r:
    input_signal_r[i*sample_rate] = 0.5
input_signal = input_signal_l + input_signal_r

# convolve inputs with IRF:    
convolved_signal = (sp.convolve(input_signal, IRF, 'full'))[:-(IRF.shape[0]-1)]

# let's add some noise:
convolved_signal_noise = convolved_signal + np.random.normal(0,0.1,len(convolved_signal))

# plot simulated convolved signal with noise:
timepoints = np.linspace(0,duration,duration*sample_rate)
fig = plt.figure()
plt.plot(timepoints, convolved_signal_noise, 'b')
plt.legend(['BOLD time series'], loc=1)
for i in times_l:
    plt.axvline(i, color='r', alpha=0.5)
for i in times_r:
    plt.axvline(i, color='g', alpha=0.5)
# plt.legend(['"true" neural signal', 'pupil time series'], loc=1)
plt.title('BOLD time series, with measurement noise')
plt.xlabel('time (s)')
plt.ylabel('a.u.')

duration = 60 # in seconds
times_l = np.array([5,15,16,25])
times_r = np.array([35,45,46,55])
regr_l = np.zeros(duration * sample_rate)
regr_r = np.zeros(duration * sample_rate)
for i in times_l:
    regr_l[i*sample_rate] = 1
for i in times_r:
    regr_r[i*sample_rate] = 1

regr_l_convolved = (sp.convolve(regr_l, IRF, 'full'))[:-(IRF.shape[0]-1)]
regr_r_convolved = (sp.convolve(regr_r, IRF, 'full'))[:-(IRF.shape[0]-1)]

# demean measured data and regressors:
regr_l_convolved = regr_l_convolved - regr_l_convolved.mean()
regr_r_convolved = regr_r_convolved - regr_r_convolved.mean()
convolved_signal_noise = convolved_signal_noise - convolved_signal_noise.mean()

timepoints = np.linspace(0,duration,duration*sample_rate)
fig = plt.figure()
plt.plot(timepoints, convolved_signal_noise, 'b')
plt.plot(timepoints, regr_l_convolved, color='r', lw=2)
plt.plot(timepoints, regr_r_convolved+0.05, color='g', lw=2)
plt.legend(['BOLD time series', 'regressor left events', 'regressor right events'], loc=1)
plt.title('regressors')
plt.xlabel('time (s)')
plt.ylabel('a.u.')

# make design matrix:
designMatrix = np.vstack((regr_l_convolved, regr_r_convolved))

# # plot design matrix:
# fig = plt.figure(figsize=(2,2))
# plt.imshow(designMatrix.T)
# plt.xticks([0,1])
# plt.title('design matrix')
# plt.xlabel('nr samples')
# plt.ylabel('length run')

# multiple regression:
designMatrix = np.mat(designMatrix).T
betas = ((designMatrix.T * designMatrix).I * designMatrix.T) * np.mat(convolved_signal_noise).T
betas = np.array(betas)

print 'betas: {}'.format(betas.ravel())


explained_signal = regr_l_convolved*betas[0] + regr_r_convolved*betas[1]
residuals = convolved_signal_noise - explained_signal

fig = plt.figure()
plt.plot(timepoints, convolved_signal_noise, 'b')
plt.plot(timepoints, explained_signal, 'm', lw=2)
plt.legend(['BOLD timeseries', 'explained signal'], loc=1)
plt.title('predicted signal')
plt.xlabel('time (s)')
plt.ylabel('a.u.')

fig = plt.figure()
plt.plot(timepoints, residuals, 'c')
plt.legend(['residuals'], loc=1)
plt.title('residuals')
plt.xlabel('time (s)')
plt.ylabel('a.u.')

fig = plt.figure()
ax = fig.add_subplot(111)
a = ax.imshow(betas.T, interpolation='nearest', vmin=0, vmax=2, cmap='seismic') # cmap='YlOrRd')
ax.set_xticks([0,1])
ax.set_yticks([])
ax.set_xticklabels(['beta left', 'beta right'])
fig.colorbar(a, ticks=[0, 1, 2], orientation='vertical', )

sizex = 30
sizey = 30
x, y = np.mgrid[-sizex+10:sizex+10+1, -sizey+10:sizey+10+1]
g = np.exp(-0.333*(x**2/float(sizex)+y**2/float(sizey)))
PRF = g/g.sum()
plt.imshow(PRF)
plt.xticks([]); plt.yticks([])
plt.title('"Real" pRF')