import numpy as np
import matplotlib.pyplot as plt

khu = plt.imread('http://jonghank.github.io/ee786/files/khu_x_2.png')
m, n = khu.shape[0:2]

khu_r = khu[:,:,0]
khu_g = khu[:,:,1]
khu_b = khu[:,:,2]

plt.figure(figsize=(18,9))
plt.subplot(121)
plt.imshow(khu)
plt.axis('off')
plt.title('Loaded image')
plt.show()

plt.figure()
plt.imshow(khu[110:150, 20:60, :])
plt.axis('off')
plt.title('Image zoomed in')
plt.show()

# your code here

diff_rgb = np.abs(khu_r - khu_g) + np.abs(khu_g - khu_b)
m_known, n_known = np.where(diff_rgb>1e-3)

Known_pixels = np.zeros((m,n,3))

Known_pixels[m_known, n_known, 0] = khu_r[m_known, n_known]
Known_pixels[m_known, n_known, 1] = khu_g[m_known, n_known]
Known_pixels[m_known, n_known, 2] = khu_b[m_known, n_known]

plt.figure(figsize=(18,9))
plt.subplot(121)
plt.imshow(Known_pixels)
plt.axis('off')
plt.title('Pixels with full color information')
plt.show()

# your code here

khu_gray = 0.299*khu_r + 0.587*khu_g + 0.114*khu_b

plt.figure(figsize=(18,9))
plt.subplot(121)
plt.imshow(khu_gray, cmap='gray')
plt.axis('off')
plt.title('Full grayscale image')
plt.show()

# your code here

import cvxpy as cp

Xr = cp.Variable((m,n))
Xg = cp.Variable((m,n))
Xb = cp.Variable((m,n))

variables = [Xr, Xg, Xb]

obj = cp.tv(*variables) + 1000*cp.sum_squares(0.299*Xr+0.587*Xg+0.114*Xb-khu_gray)
obj = cp.Minimize(obj)
constraints = [
    Xr[m_known, n_known] == khu_r[m_known, n_known] ,
    Xg[m_known, n_known] == khu_g[m_known, n_known] , 
    Xb[m_known, n_known] == khu_b[m_known, n_known] 
]

prob = cp.Problem(obj, constraints)
prob.solve(verbose=True, solver=cp.SCS)

Xrec = np.zeros((m, n, 3))
Xrec[:,:,0] = np.clip(Xr.value, 0, 1)
Xrec[:,:,1] = np.clip(Xg.value, 0, 1)
Xrec[:,:,2] = np.clip(Xb.value, 0, 1)

plt.figure(figsize=(18,9))
plt.subplot(121)
plt.imshow(khu)
plt.title('Grayscale with 2% color pixels')
plt.axis('off')
plt.show()
plt.figure(figsize=(18,9))
plt.subplot(121)
plt.imshow(Xrec)
plt.title('Colorized')
plt.axis('off')
plt.show()

# your code here

khu_color = plt.imread('http://jonghank.github.io/ee786/files/khu_entrance.png')

plt.figure(figsize=(18,9))
plt.subplot(121)
plt.imshow(khu)
plt.title('Grayscale with 2% color pixels')
plt.axis('off')
plt.subplot(122)
plt.imshow(Known_pixels)
plt.axis('off')
plt.title('Pixels with full color information')
plt.show()
plt.figure(figsize=(18,9))
plt.subplot(121)
plt.imshow(Xrec)
plt.title('Colorized')
plt.axis('off')
plt.subplot(122)
plt.imshow(khu_color)
plt.title('Original color image')
plt.axis('off')
plt.show()