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

import numpy as np
from copy import deepcopy

class PCA():
    '''
    Formula:
        X = U . S . Vt
    Useage:
        p = PCA(X, fraction = 0.9)
        in which `X` is the matrix of your data and `fraction` means
        use principal components that account for e.g. 0.9 of the
        total variance
    Out:
        p.U, p.S, p.Vt from numpy.linalg.svd
        p.eigen: the eigenvalues of A*A, in decreasing order
            eigen[j] / eigen.sum() is variable j's fraction of the
            total variance;
            look at the first few eigen[] to see how many PCs get
            to 90 %, 95 % ...
        p.npc: number of principal components,
            e.g. 2 if the top 2 eigenvalues are >= `fraction` of the total.
        p.sumvariance: the number of the variances
    Methods:
        p.center(X, axis = 0, scale = True)
            classmethod for centering the matrix (e.g. `X`), returns
            the centered matrix with out changing the original matrix
    '''
    def __init__(self, X, fraction = 0.85):
        assert 0 <= fraction <= 1
        # center the matric
        A = self.center(X)
        # SVD
        self.U, self.S, self.Vt = np.linalg.svd(A, full_matrices = False)
        assert np.all(self.S[:-1] >= self.S[1:]) # sorted
        self.eigen = self.S**2
        self.sumvariance = np.cumsum(self.eigen)
        self.sumvariance /= self.sumvariance[-1]
        self.npc = np.searchsorted(self.sumvariance, fraction) + 1
        self.pc = self.U[:, :self.npc] * self.S[:self.npc]
    @classmethod
    def center(self, X, axis = 0, scale = True):
        A = deepcopy(X)
        mean = A.mean(axis = axis)
        A -= mean
        if scale:
            std = A.std(axis = axis)
            A /= np.where(std, std, 1.0)
        return A

if __name__ == '__main__':
    from matplotlib import pyplot as plt
    import sys
    if len(sys.argv) >= 2:
        csv = sys.argv[1]
    else:
        print '%s /path/to/your_data_file_in_csv_format' %sys.argv[0]
        sys.exit(1)
    # you can define your own colors here
    colors = {(0, 1, 2): 'red',
            (3, 4, 5): 'blue',
            (6, 7, 8, 9): 'green',
            (10, 11, 12): 'yellow'}
    def get_color(idx):
        for i, c in colors.items():
            if idx in i:
                return c
    X = np.genfromtxt(csv, delimiter = ',')
    p = PCA(X, fraction = 0.9)
    print p.npc
    pc = p.pc
    x, y = pc[:, 0], pc[:, 1]
    fig = plt.figure(figsize=(8,6))
    ax = fig.add_subplot(111)
    for i, k in enumerate(zip(x,y)):
        ax.text(k[0], k[1], str(i), bbox=dict(facecolor = get_color(i), alpha=0.5))
    #plt.xlim(-8, 8)
    #plt.ylim(-4, 4)
    plt.xlabel('PC1')
    plt.ylabel('PC2')
    plt.title('bulabula')
    plt.legend()
    plt.grid(True)
    plt.show()