import rpca
from rpca import data

# Load "Sleep in Mammals" database
X = rpca.data.load_sleep()

# Transform it using Robust PCA
huber_loss = rpca.loss.HuberLoss(delta=1)
rpca_transformer = rpca.MRobustPCA(2, huber_loss)
X_rpca = rpca_transformer.fit_transform(X)

import matplotlib.pyplot as plt
%matplotlib inline

# Plot progress during iterations
plt.figure(figsize=(8, 5))
plt.plot(range(1, rpca_transformer.n_iterations_ + 1), rpca_transformer.errors_, '-o')
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.show()

# Compare our method with regular PCA
from sklearn.decomposition import PCA
pca_transformer = PCA(2)
X_pca = pca_transformer.fit_transform(X)

# Calculate error norms for both methods
import numpy as np
def error_norms(components):
    n_samples, n_features = X.shape
    non_projected_matrix = np.identity(n_features) - components.T.dot(components)
    return np.sqrt(np.diag(X.dot(non_projected_matrix.dot(X.T))))

errors_pca, errors_rpca = error_norms(pca_transformer.components_), error_norms(rpca_transformer.components_)

plt.figure(figsize=(8, 5))
plt.plot(sorted(errors_pca, key=lambda x:-x)[:20], '-o', label='PCA', color='r')
plt.plot(sorted(errors_rpca, key=lambda x:-x)[:20], '-o', label='Robust PCA', color='b')
plt.title('Sorted error norms')
plt.legend()
plt.show()