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.')

# input:
duration = 35 # in seconds
times = np.array([1,6,11,11.5,16,16.5,21,22,26,27])
input_signal = np.zeros(duration * sample_rate)
for i in times:
    input_signal[i*sample_rate] = 1

# 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.25,len(convolved_signal))

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

# times for epoching:
epoch_times = [0, 3] # in seconds

# mean response:
epochs = np.vstack([convolved_signal_noise[i+(epoch_times[0]*sample_rate):i+(epoch_times[1]*sample_rate)] for i in times*sample_rate])
mean_response = np.mean(epochs, axis=0)

# plot mean response versus IRF:
timepoints = np.linspace(0,3,3*sample_rate)
fig = plt.figure(figsize=(6,6))
timepoints = np.linspace(epoch_times[0],epoch_times[1],(epoch_times[1]-epoch_times[0])*sample_rate)
fig.add_subplot(211)
for data in epochs:
    plt.plot(timepoints,data, color='b')
plt.title('epoched responses')
plt.xlabel('time from event (s)')
plt.ylabel('a.u.')
fig.add_subplot(212)
plt.plot(timepoints, mean_response, color='b')
plt.plot(timepoints, IRF, color='r')
plt.legend(['mean epoched response', 'true response'])
plt.title('mean response')
plt.xlabel('time from event (s)')
plt.ylabel('a.u.')
fig.tight_layout()

# make design matrix:
nr_samples = 3 * sample_rate # here we define the length of the deconvolution response we're interested in (30 samples = 3-s in our case).
designMatrix = np.zeros((nr_samples, duration*sample_rate))
for i in (times*sample_rate):
    for j in range(int(nr_samples)):
        designMatrix[j,i+j] = 1

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

# deconvolution:
designMatrix = np.mat(designMatrix).T
deconvolved_response = ((designMatrix.T * designMatrix).I * designMatrix.T) * np.mat(convolved_signal_noise).T
deconvolved_response = np.array(deconvolved_response)

# plot deconvoled response versus true response:
timepoints = np.linspace(0,3,3*sample_rate)
fig = plt.figure(figsize=(6,6))
fig.add_subplot(211)
plt.plot(timepoints,deconvolved_response, color='b')
plt.xlim(xmax=3)
plt.ylim(ymin=-0.5, ymax=3.5)
plt.legend(['deconvolved response'])
plt.title('deconvolved response')
plt.xlabel('time (s)')
plt.ylabel('a.u.')
fig.add_subplot(212)
plt.plot(timepoints,IRF, color='r')
plt.xlim(xmax=3)
plt.ylim(ymin=-0.5, ymax=3.5)
plt.legend(['true response'])
plt.xlabel('time from event (s)')
plt.ylabel('a.u.')