import numpy as np
import matplotlib.pyplot as plt
import copy

def doPCA(x):
    from sklearn.preprocessing import StandardScaler
    from sklearn.decomposition import PCA
    x = StandardScaler().fit_transform(x)
    pca = PCA(n_components=2) # How many components
    pcComponents = pca.fit_transform(x)
    return pcComponents

def calculateGaussian(x):
    from scipy import stats
    X, Y = np.mgrid[-1.25:1.25:100j, -1.25:1.25:100j]
    positions = np.vstack([X.ravel(), Y.ravel()])
    kernel = stats.gaussian_kde(x.T)

    Z = np.reshape(kernel(positions).T, X.shape)
    return np.rot90(Z)

def minmaxNormalize(x):
    m = x.min()
    M = x.max()
    return (x - m) / (M - m) 

def calculateRGB(rep, label):
    # Reduce Dimension by PCA
    pc = doPCA(copy.deepcopy(rep))

    # Normalizing Representation
    from sklearn.preprocessing import normalize
    pc = normalize(pc)

    # Calculate Value
    V = calculateGaussian(pc)
    V = minmaxNormalize(V) * 255

    # Calculate and Smoothing Hue
    from scipy import interpolate
    X, Y = np.mgrid[-1.25:1.25:100j, -1.25:1.25:100j]
    H = interpolate.griddata(pc, label, (X, Y), method='nearest')
    H = (minmaxNormalize(H) * 150).astype(np.uint8)
    import cv2
    H = cv2.GaussianBlur(H,(5,5),0)

    S = np.ones_like(V) * 255

    # Convert Color Space
    hsv = np.dstack([H, V, S]).astype(np.uint8)
    return cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)

def calc_label_distance(rep, label):
    maximum_mean = 1e9
    MAX_ITER = 50
    ind = 0
    for k in range(MAX_ITER):
        label_distance = np.sum((label - label[k]) ** 2, axis=2).mean(axis=1)
        mean = np.mean(label_distance)
        if mean < maximum_mean:
            maximum_mean = mean
            ind = k
    label_distance = np.sum((label - label[ind]) ** 2, axis=2).mean(axis=1)
    rep = rep[MAX_ITER:]
    label_distance = np.log(label_distance[MAX_ITER:])
    return rep, label_distance

def remove_outlier(rep, label, deviation=2.5):
    # Remove outlier Start
    mean = np.mean(label)
    standard_deviation = np.std(label)

    distance_from_mean = abs(label - mean)
    max_deviations = deviation
    not_outlier = distance_from_mean < max_deviations * standard_deviation

    rep = rep[not_outlier]
    label = label[not_outlier]
    return rep, label

# rep: representation
# label: label
rep, label = calc('regcon_linear_out.ckpt')
rep, label = calc_label_distance(rep, label)
rep, label = remove_outlier(rep, label)
rgb = calculateRGB(rep, label)

# Draw Figure
fig, ax = plt.subplots(1, 1, figsize=(3,3))
ax.axes.set_aspect(aspect=1)
ax.imshow(rgb, extent=[-1.25, 1.25, -1.25, 1.25])        
fig.set_size_inches(4, 4)
ax.set_xlim([-1.25, 1.25])
ax.set_ylim([-1.25, 1.25])
ax.set_title('Alignment & Uniformity')

plt.show()

