'''
lecture21_pca_cov.py
PCA dimension reduction on and plotting
Oliver W. Layton
CS251: Data Analysis and Visualization
Spring 2019
'''
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_style("darkgrid")

data = sns.load_dataset('diamonds')
data = data.sample(1000)
print(data[:3])

# Assign data variables x (price),y (carat),z (z) and n
x = np.matrix(data['price']).T
y = np.matrix(data['carat']).T
z = np.matrix(data['z']).T
n = x.size

# 1) Center data: xc, yc, zc
xc = x - np.mean(x)
yc = y - np.mean(y)
zc = z - np.mean(z)

# Make data matrix
dataMat = np.hstack([xc, yc, zc])
print(dataMat.shape)

# Compute data ranges of each variable
diff = np.max(dataMat, axis=0) - np.min(dataMat, axis=0)

# 2) Normalize so range is 1
dataMat = dataMat / diff
# print(dataMat)
print('Centered variable mins:\n', np.min(dataMat, axis=0))
print('Centered variable maxs:\n', np.max(dataMat, axis=0))

# 3) Compute covariance matrix
sigma = (dataMat.T * dataMat) / (n-1)
print('Covariance matrix:\n', sigma)


# 4) Compute eigenvalues and eigenvectors
eVals, eVecs = np.linalg.eig(sigma)
print('Eigenvalues:\n', eVals)
print('Eigenvectors:\n', eVecs)

# 5) Sort eigenvalues (already done)

# 6) Compute portion of variance accounted for by each eigenvalue
propVar = [eVal / np.sum(eVals) for eVal in eVals]
print('Proportion variance accounted for:', propVar)

# 7) Remove eigenvectors based on eigenvalue portion variance scores (Step 6)
# eVals = eVals
# eVecs = eVecs
eVals = eVals[:2]
eVecs = eVecs[:, :2]

# 8) Project into PCA space.
# Convert to ndarray for plotting
dataPCAProj = np.array(dataMat * eVecs)

# 8a) If whitening, divide by squart of eigenvalues
# dataPCAProj = np.array(dataMat * eVecs) @ np.diag(1/np.sqrt(eVals))
# print(np.diag(1/np.sqrt(eVals)))
# print('New cov mat:\n', np.cov(dataPCAProj, rowvar=False))


# Plot data in original space
xCarat = data['carat']
yPrice = data['price']
zZ = data['z']
plt.scatter(xCarat, yPrice)
plt.xlabel('carat')
plt.ylabel('price')
# sns.scatterplot(x='carat', y='price', data=data)
plt.show()
plt.scatter(zZ, yPrice)
plt.xlabel('z')
plt.ylabel('price')
# sns.scatterplot(x='z', y='price', data=data)
plt.show()

# Plot data in PCA space
plt.scatter(dataPCAProj[:, 0].flatten(), dataPCAProj[:, 1].flatten())
plt.title('data in PCA space')
plt.xlabel('PC1')
plt.ylabel('PC2')
plt.show()

# Price dominates eVals so we conclude that PC0 mostly corresponds to price.
# It accounts for most of the variance out of (price, carat, z) vars.

# 9) (optional) Reconstruct data based on top PCs.
# In this case, we have to comment out the normalization step above (#2)
# otherwise the shape of the reconstructed data will look ok, but at unit scale.
origMeans = np.hstack([np.mean(x), np.mean(y), np.mean(z)])
data2 = np.matrix(dataPCAProj) * eVecs.T + origMeans
print(data2.shape)

plt.scatter([data2[:, 1]], [data2[:, 0]])
plt.title('Reconstructed data')
plt.xlabel('carat')
plt.ylabel('price')
plt.show()