import helper

def generate_signal(num_steps, epsilon):
    return helper.gaussian_spectrum_1D(np.random.rand(num_steps), epsilon)

plt.plot(generate_signal(1000, 1.))
plt.plot(generate_signal(1000, 0.0001))

"""
Explanation: Exercise 1
1. Function that generates random signal
End of explanation
"""


T = 1000
x1 = generate_signal(T, 1 / 30**2)
x2 = generate_signal(T, 1 / 100**2)

x1 = (x1 - np.mean(x1)) / np.std(x1)
x2 = (x2 - np.mean(x2)) / np.std(x2)
np.mean(x1), np.mean(x2), np.std(x1), np.std(x2)

plt.plot(x1, label=r'$\tau$ = 30')
plt.plot(x2, label=r'$\tau$ = 100')
plt.xlabel('Step')
plt.ylabel('Signal Value')
plt.legend()

"""
Explanation: 2. Generate two signals
End of explanation
"""


X = np.vstack((x1, x2))
X.shape

C = np.cov(X)
C

X_dot = np.diff(X)
C_dot = np.dot(X_dot, X_dot.T) / (T - 1)
C_dot

from scipy.linalg import eig

eigenvalues, eigenvectors = eig(C_dot, C)
for i in range(len(eigenvalues)):
    eigenvalue = eigenvalues[i]
    normalized_eigenvector = eigenvectors[:, i] / np.linalg.norm(eigenvectors[:, i])
    print('Eigenvalue: {} - Normalized Eigenvector: {}'.format(eigenvalue, normalized_eigenvector))

"""
Explanation: The signal with higher $\tau$ is smoother, i.e. it contains less high-frequency components.
3. Linear SFA
End of explanation
"""


Ts = np.logspace(1, 4, 20).astype(int)
eigenvector_values = []

for T in Ts:
    x1 = generate_signal(T, 1 / 30**2)
    x2 = generate_signal(T, 1 / 100**2)
    
    x1 = (x1 - np.mean(x1)) / np.std(x1)
    x2 = (x2 - np.mean(x2)) / np.std(x2)

    X = np.vstack((x1, x2))
    C = np.cov(X)
    X_dot = np.diff(X)
    C_dot = np.dot(X_dot, X_dot.T) / (T - 1)
    
    eigenvalues, eigenvectors = eig(C_dot, C)
    ev1 = eigenvectors[:, 0] / np.linalg.norm(eigenvectors[:, 0])
    ev2 = eigenvectors[:, 1] / np.linalg.norm(eigenvectors[:, 1])
    
    eigenvector_values.append(np.hstack((ev1, ev2)))
    
eigenvector_values = np.array(eigenvector_values)

plt.plot(Ts, eigenvector_values, 'o')
plt.xscale('log')
plt.xlabel('Sample size')
plt.ylabel('Eigenvector values')

"""
Explanation: 4. Mixing
The eigenvector with the smallest eigenvalue corresponds to the slowest SFA signal. Here, this eigenvector is almost (0, 1), i.e. the SFA signal contains almost only the second input signal ($x_2$). This makes sense because this is the slower input signal (higher $\tau$ value). Accordingly, the second SFA signal contains almost only the other input signal ($x_1$).
5. Sample size
End of explanation
"""


signal = generate_signal(10000, 0.04)
plt.plot(signal)

"""
Explanation: While three components of the two eigenvectors are stable with regards to the sample size, one of them is unstable for a small sample size and stabilizes around 0 for longer input signals. This may have to do with the fact, that low-frequency signals are not apparent if the sample size is too small.
Exercise 2
1. Generate image
End of explanation
"""


X = np.zeros((100, 9900))

for i in range(9900):
    X[:, i] = signal[i:i+100]

plt.plot(X[:, :3])

plt.plot(X[0])
plt.xlabel('Time')
plt.ylabel('Activation of first visual cell')

"""
Explanation: 2. Receptive field input
End of explanation
"""


C = np.cov(X)
C

X_dot = np.diff(X)
C_dot = np.dot(X_dot, X_dot.T) / (T - 1)
C_dot

eigenvalues, eigenvectors = eig(C_dot, C)

"""
Explanation: 3. SFA
End of explanation
"""


indices_for_sorting = np.argsort(eigenvalues)
sorted_eigenvectors = eigenvectors.T[indices_for_sorting].T

for i in range(5):
    plt.plot(eigenvectors[:, i], 'o', label='Eigenvector {}'.format(i))
plt.xlabel('Component of Eigenvector')
plt.ylabel('Value')
plt.legend()

"""
Explanation: 4. Sorting
End of explanation
"""


Y = np.dot(eigenvectors.T, X)
Y.shape

for i in range(5):
    plt.plot(Y[i], label=i)
plt.xlabel('Time')
plt.ylabel('SFA signal $y_i$')
plt.legend()

"""
Explanation: 5. Discussion
The values of the eigenvector components have a sinusoidal shape. Eigenvectors corresponding to slower SFA components have lower frequency. 
6. Projection on SFA-basis
End of explanation
"""


