#!/usr/bin/env python
# -*- coding: utf-8 -*-


import numpy as np
import scipy as sp
import matplotlib.pyplot as plt


def make_data(n_samples=1000):
    m = np.asarray([1, 1])
    cov = np.asarray([[10, 3], [3, 2]])
    # cov = np.asarray([[20, 2], [2, 5]])
    # cov = np.asarray([[20, 0], [0, 2]])
    # m = np.asarray([1, 1, 1])
    # cov = np.asarray([[10, 3, 0], [3, 2, 0], [0, 0, 1]])
    data = sp.random.multivariate_normal(m, cov, size=n_samples)
    return data


def test_pca():
    import pca
    X = make_data(1000)

    PCA = pca.PCA(n_components=2)
    PCA.fit(X)
    p1 = PCA.mean + np.sqrt(PCA.variances[0]) * PCA.components[0, :]
    p2 = PCA.mean + np.sqrt(PCA.variances[1]) * PCA.components[1, :]
    line1 = np.vstack((PCA.mean, p1))  # Line of principal component from mean to p1
    line2 = np.vstack((PCA.mean, p2))  # Line of principal component from mean to p2

    PCA2 = pca.PCA(n_components=1)
    PCA2.fit(X)
    X_clean = PCA2.reconstruct(PCA2.project(X))  # Reconstruct from projected data.

    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111)
    ax.set_title("PCA example")
    ax.scatter(X[:, 0], X[:, 1], alpha=0.5, label="Noisy")
    ax.scatter(X_clean[:, 0], X_clean[:, 1], c='y', alpha=0.5, label="Clean")
    ax.plot(line1[:, 0], line1[:, 1], 'c-', linewidth=3, label="Principal component 1")
    ax.plot(line2[:, 0], line2[:, 1], 'g-', linewidth=3, label="Principal component 2")
    ax.legend()

    rmin = np.min(X) + 1
    rmax = np.max(X) + 1
    plt.xlim(rmin, rmax)  # Adjust the axis limits else plot will be distorted.
    plt.ylim(rmin, rmax)
    plt.show()


if __name__ == '__main__':
    test_pca()
