from pylab import eig, argsort, mean, tile, newaxis, loadtxt, dot, \
     show, plot, figure, arange, legend

def sortedeig(M):
    d, U = eig(M)
    si = argsort(d)[-1::-1]
    return (d[si], U[:,si])

def calc_sortedeig(**kwargs):
    """Calculate the sorted eigenvalues and eigenvectors of the data set."""
    if kwargs['data'] == 'natural':
        X = loadtxt('natural400_700_5.asc').T
        N = 219
    elif kwargs['data'] == 'munsell':
        X = loadtxt('munsell380_800_1.asc').reshape(1269, 421).T
        N = 1269
    else:
        raise ValueError('Undefined data set.')

    Xzm = X - tile(mean(X, 1)[:,newaxis], N)
    S = dot(Xzm, Xzm.T) / (N - 1)
    return sortedeig(S)

def PCA(**kwargs):
    """Show scree diagram of a data set."""
    d, U = calc_sortedeig(**kwargs)
    figure(1)
    plot(d)
    show()

def EigenImages(k, **kwargs):
    """Plot the first k eigenvectors of the data set."""
    d, U = calc_sortedeig(**kwargs)
    if kwargs['data'] == 'natural':
        min, max, step = 400, 701, 5
    elif kwargs['data'] == 'munsell':
        min, max, step = 380, 801, 1
    figure(2)
    a = arange(min, max, step)
    for i in range(k):
        plot(a, U[:,i], label=str(i + 1))
    legend()
    show()

def Reconstruct(k, sample, **kwargs):
    """Reconstruct the original spectrum from the k principle components."""
    d, U = calc_sortedeig(**kwargs)
    if kwargs['data'] == 'natural':
        X = loadtxt('natural400_700_5.asc').T
        min, max, step = 400, 701, 5
    elif kwargs['data'] == 'munsell':
        X = loadtxt('munsell380_800_1.asc').reshape(1269, 421).T
        min, max, step = 380, 801, 1

    # Select the specified vector, subtract the mean from it and multiply with
    # the transposed eigenvector basis to get the coordinates with respect to U.
    # Then, take the first k components and try to reconstruct the original
    # spectrum.
    x = X[:,sample]
    xbar = mean(X, 1)
    yzm = dot(U.T, x - xbar)[:k]
    x_k = dot(U[:,:k], yzm[:k]) + xbar

    a = arange(min, max, step)
    figure(3)
    plot(a, x, label='original')
    plot(a, x_k, label='reconstructed')
    legend()
    show()

if __name__ == '__main__':
    PCA(data='natural')
    PCA(data='munsell')

    EigenImages(5, data='natural')

    Reconstruct(5, 23, data='natural')
