"""ICA Example 1: Mix and separate two sources using FastICA.

Python equivalent of matlabex1.m
Requires: numpy, matplotlib, scikit-learn
"""
import numpy as np
import matplotlib.pyplot as plt
from sklearn.decomposition import FastICA

# Define two independent sources
t = np.linspace(0, 50, 1000)
A = np.sin(t)
B = np.sin(np.linspace(0, 37, 1000) + 5)

fig, axes = plt.subplots(2, 1)
axes[0].plot(A)
axes[0].set_title("Source A")
axes[1].plot(B, "r")
axes[1].set_title("Source B")
plt.tight_layout()
plt.show()

# Mix the two sources linearly
M1 = A - 2 * B
M2 = 1.73 * A + 3.41 * B

fig, axes = plt.subplots(2, 1)
axes[0].plot(M1)
axes[0].set_title("Mixture 1 = A − 2B")
axes[1].plot(M2, "r")
axes[1].set_title("Mixture 2 = 1.73A + 3.41B")
plt.tight_layout()
plt.show()

# Unmix using FastICA
ica = FastICA(n_components=2, random_state=0)
sources = ica.fit_transform(np.column_stack([M1, M2]))

fig, axes = plt.subplots(1, 2, figsize=(10, 3))
axes[0].plot(sources[:, 0])
axes[0].set_title("Recovered source 1")
axes[1].plot(sources[:, 1])
axes[1].set_title("Recovered source 2")
plt.tight_layout()
plt.show()