# image processing
from PIL import Image
from io import BytesIO
import webcolors

# data analysis
import math
import numpy as np
import pandas as pd

# visualization
import matplotlib.pyplot as plt
from importlib import reload
from mpl_toolkits import mplot3d
import seaborn as sns

# modeling
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.preprocessing import MinMaxScaler

ori_img = Image.open("images/lena.png")
ori_img

X = np.array(ori_img.getdata())
ori_pixels = X.reshape(*ori_img.size, -1)
ori_pixels.shape

def imageByteSize(img):
  img_file = BytesIO()
  image = Image.fromarray(np.uint8(img))
  image.save(img_file, 'png')
  return img_file.tell()/1024
ori_img_size = imageByteSize(ori_img)
ori_img_n_colors = len(set(ori_img.getdata()))

ori_img_total_variance = sum(np.linalg.norm(X - np.mean(X, axis = 0), axis = 1)**2)

kmeans = KMeans(n_clusters = 2,
        n_jobs = -1,
        random_state = 123).fit(X)
kmeans_df = pd.DataFrame(kmeans.cluster_centers_, columns = ['Red', 'Green', 'Blue'])

def closest_colour(requested_colour):
  min_colours = {}
  for key, name in webcolors.CSS3_HEX_TO_NAMES.items():
    r_c, g_c, b_c = webcolors.hex_to_rgb(key)
    rd = (r_c - requested_colour[0]) ** 2
    gd = (g_c - requested_colour[1]) ** 2
    bd = (b_c - requested_colour[2]) ** 2
    min_colours[(rd + gd + bd)] = name
  return min_colours[min(min_colours.keys())]
def get_colour_name(requested_colour):
  try:
    closest_name = actual_name = webcolors.rgb_to_name(requested_colour)
  except ValueError:
    closest_name = closest_colour(requested_colour)
  return closest_name
kmeans_df["Color Name"] = list(map(get_colour_name, np.uint8(kmeans.cluster_centers_)))
kmeans_df

def replaceWithCentroid(kmeans):
  new_pixels = []
  for label in kmeans.labels_:
    pixel_as_centroid = list(kmeans.cluster_centers_[label])
    new_pixels.append(pixel_as_centroid)
  new_pixels = np.array(new_pixels).reshape(*ori_img.size, -1)
  return new_pixels
new_pixels = replaceWithCentroid(kmeans)

def plotImage(img_array, size):
  reload(plt)
  plt.imshow(np.array(img_array/255).reshape(*size))
  plt.axis('off')
  return plt
  
plotImage(new_pixels, new_pixels.shape).show()

def calculateBCSS(X, kmeans):
  _, label_counts = np.unique(kmeans.labels_, return_counts = True)
  diff_cluster_sq = np.linalg.norm(kmeans.cluster_centers_ - np.mean(X, axis = 0), axis = 1)**2
  return sum(label_counts * diff_cluster_sq)
WCSS = kmeans.inertia_
BCSS = calculateBCSS(X, kmeans)
exp_var = 100*BCSS/(WCSS + BCSS)
print("WCSS: {}".format(WCSS))
print("BCSS: {}".format(BCSS))
print("Explained Variance: {:.3f}%".format(exp_var))
print("Image Size: {:.3f} KB".format(imageByteSize(new_pixels)))

range_k_clusters = (2, 21)
kmeans_result = []
for k in range(*range_k_clusters):
  # CLUSTERING
  kmeans = KMeans(n_clusters = k,
          n_jobs = -1,
          random_state = 123).fit(X)
  
  # REPLACE PIXELS WITH ITS CENTROID
  new_pixels = replaceWithCentroid(kmeans)
  
  # EVALUATE
  WCSS = kmeans.inertia_
  BCSS = calculateBCSS(X, kmeans)
  exp_var = 100*BCSS/(WCSS + BCSS)
  
  metric = {
    "No. of Colors": k,
    "Centroids": list(map(get_colour_name, np.uint8(kmeans.cluster_centers_))),
    "Pixels": new_pixels,
    "WCSS": WCSS,
    "BCSS": BCSS,
    "Explained Variance": exp_var,
    "Image Size (KB)": imageByteSize(new_pixels)
  }
  
  kmeans_result.append(metric)
kmeans_result = pd.DataFrame(kmeans_result).set_index("No. of Colors")
kmeans_result

def locateOptimalElbow(x, y):
  # START AND FINAL POINTS
  p1 = (x[0], y[0])
  p2 = (x[-1], y[-1])
  
  # EQUATION OF LINE: y = mx + c
  m = (p2[1] - p1[1]) / (p2[0] - p1[0])
  c = (p2[1] - (m * p2[0]))
  
  # DISTANCE FROM EACH POINTS TO LINE mx - y + c = 0
  a, b = m, -1
  dist = np.array([abs(a*x0+b*y0+c)/math.sqrt(a**2+b**2) for x0, y0 in zip(x,y)])
  return np.argmax(dist) + x[0]

def calculateDerivative(data):
  derivative = []
  for i in range(len(data)):
    if i == 0:
      # FORWARD DIFFERENCE
      d = data[i+1] - data[i]
    elif i == len(data) - 1:
      # BACKWARD DIFFERENCE
      d = data[i] - data[i-1]
    else:
      # CENTER DIFFERENCE
      d = (data[i+1] - data[i-1])/2
    derivative.append(d)
  return np.array(derivative)
def locateDrasticChange(x, y):
  # CALCULATE GRADIENT BY FIRST DERIVATIVE
  first_derivative = calculateDerivative(np.array(y))
  
  # CALCULATE CHANGE OF GRADIENT BY SECOND DERIVATIVE
  second_derivative = calculateDerivative(first_derivative)
return np.argmax(np.abs(second_derivative)) + x[0]

optimal_k = []
for col in kmeans_result.columns[2:]:
  optimal_k_dict = {}
  optimal_k_dict["Metric"] = col
  if col == "Image Size (KB)":
    optimal_k_dict["Method"] = "Derivative"
    optimal_k_dict["Optimal k"] = locateDrasticChange(kmeans_result.index, kmeans_result[col].values)
  else:
    optimal_k_dict["Method"] = "Elbow"
    optimal_k_dict["Optimal k"] = locateOptimalElbow(kmeans_result.index, kmeans_result[col].values)
  optimal_k.append(optimal_k_dict)
optimal_k = pd.DataFrame(optimal_k)
optimal_k

relative_size = ori_vs_kmeans.loc["Color-Reduced", "Image Size (KB)"]/ori_vs_kmeans.loc["Original", "Image Size (KB)"]
print("Reduction: {:.3f}% from original image size".format((1-relative_size)*100))
print("Explained Variance: {:.3f}%".format(ori_vs_kmeans.loc["Color-Reduced", "Explained Variance"]))

res = []
cum_var = []
X_t = np.transpose(X)
for channel in range(3):
  # SEPARATE EACH RGB CHANNEL
  pixel = X_t[channel].reshape(*ori_pixels.shape[:2])
  
  # PCA
  pca = PCA(random_state = 123)
  pixel_pca = pca.fit_transform(pixel)
  
  pca_dict = {
    "Projection": pixel_pca,
    "Components": pca.components_,
    "Mean": pca.mean_
  }
  res.append(pca_dict)
  
  # EVALUATION
  cum_var.append(np.cumsum(pca.explained_variance_ratio_))

cum_var_df = pd.DataFrame(np.array(cum_var).T * 100, 
             index = range(1, pca.n_components_+1),
             columns = ["Explained Variance by Red",
                   "Explained Variance by Green",
                   "Explained Variance by Blue"])
cum_var_df["Explained Variance"] = cum_var_df.mean(axis = 1)
cum_var_df

pca_results = []
for n in range(1, pca.n_components_+1):
  # SELECT N-COMPONENTS FROM PC
  temp_res = []
  for channel in range(3):
    pca_channel = res[channel]
    pca_pixel = pca_channel["Projection"][:, :n]
    pca_comp = pca_channel["Components"][:n, :]
    pca_mean = pca_channel["Mean"]
    compressed_pixel = np.dot(pca_pixel, pca_comp) + pca_mean
    temp_res.append(compressed_pixel.T)
  compressed_image = np.transpose(temp_res)
  
  pca_dict = {
    "n": n,
    "Pixels": compressed_image,
    "Explained Variance": cum_var_df["Explained Variance"][n],
    "Image Size (KB)": imageByteSize(compressed_image),
    "No. of Colors": len(np.unique(np.uint8(compressed_image).reshape(-1, 3), axis = 0))
  }
  
  pca_results.append(pca_dict)
pca_results = pd.DataFrame(pca_results).set_index("n")
pca_results.head()

relative_size = ori_vs_pca.loc["PC-Reduced", "Image Size (KB)"]/ori_vs_kmeans.loc["Original", "Image Size (KB)"]
print("Reduction: {:.3f}% from original image size".format((1-relative_size)*100))
print("Explained Variance: {:.3f}%".format(ori_vs_pca.loc["PC-Reduced", "Explained Variance"]))

reduction_kmeans = (1-final_compare.loc["Color-Reduced", "Image Size (KB)"] / ori_img_size) * 100
reduction_pca = (1-final_compare.loc["PC-Reduced", "Image Size (KB)"] / ori_img_size) * 100
print("Image Size Reduction using K-Means: {:.3f}%".format(reduction_kmeans))
print("Image Size Reduction using PCA: {:.3f}%".format(reduction_pca))