from sys import argv, exit
from q21_multivariate import mu, cov, dataset
from pylab import array, mean, tile, newaxis, dot, eigvals, \
        axis, figure, clf, show, plot, sum

if len(argv) == 3:
    step = int(argv[2])
elif len(argv) == 2:
    step = 100
else:
    print 'Usage: python %s SAMPLES [ STEP_SIZE ]' % (argv[0])
    exit()

# `samples' is the size of the generated dataset.
samples = int(argv[1])
Y = dataset(samples)

def estimate(n):
    """Return eigenvalues of unbiased estimators for the covariance matrix
    Sigma (based on a pseudo-random generated dataset)."""
    # Sigma = 1 / (n - 1) * Sum for i=1 to n: (x_i - x_mean) T(x_i - x_mean),
    # where T(x) is the transpose of `x'. Mu = x_mean and
    # Yzm = Sum for i=1 to n: x_i - x_mean.
    sliced = [Y[i][:n] for i in range(len(Y))]
    est_mu = mean(sliced, 1)
    Yzm = sliced - tile(est_mu[:,newaxis], n)
    est_cov = dot(Yzm, Yzm.T) / (n - 1)
    return (est_mu, est_cov)

figure(1)
clf()

# Part 1 - Estimate mu and cov, experiment with various sizes of datasets.
# We use steps of 100 for the number of used samples.
X = range(step, samples + 1, step)
diff_mu = []
diff_cov = []
estimated_mu = []
for n in X:
    est_mu, est_cov = estimate(n)
    diff_mu.append(abs(sum(est_mu - mu)))
    diff_cov.append(abs(sum(est_cov - cov)))
    estimated_mu.append(est_mu)
plot(X, diff_mu)
plot(X, diff_cov)
# Observe in the following graph that the accuracy increases when more
# vectors from the generated dataset are used. There are some fluctuations due
# to the randomness of the dataset, but the difference lines should converge to
# zero.
show()

# Part 2 - Calculate covariance of estimated mu.
sample_count = len(estimated_mu)
estimated_mu = array(estimated_mu).T
m = mean(estimated_mu, 1)
Yzm = estimated_mu - tile(m[:,newaxis], sample_count)
S = dot(Yzm, Yzm.T) / (samples - 1)
# When many samples are used to calculate the estimators, the estimators will
# converge to the original values. The more samples are uses each time, the more
# convergence will occur, which will lead to smaller eigenvalues of the
# covariance matrix (printed below). For example, call the program with 1000 and
# 10000 samples and you will see that the eigenvalues will be smaller with 10000
# samples.
print eigvals(S)
