import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import requests
X = Image.open(requests.get("https://jonghank.github.io/ase7030/files/aerospace_building_corrupted.png", stream=True).raw)
X = np.array(X)/255
m, n = X.shape

plt.figure(figsize=(16,9), dpi=100)
plt.subplot(121)
plt.imshow(X, cmap='gray')
plt.axis('off')
plt.title('Corrupted image')
plt.show()

B = (X<0.9)*1

plt.figure(figsize=(16,9), dpi=100)
plt.subplot(121)
plt.imshow(X, cmap='gray')
plt.axis('off')
plt.title('Corrupted image')
plt.subplot(122)
plt.imshow(B*X, cmap='gray')
plt.axis('off')
plt.title('Known pixels')
plt.show()

import numpy as np
import scipy.sparse as ssp

mn = m*n
p = np.sum(B)
q = mn - p

B_k = B.flatten('F')
B_u = 1-B.flatten('F')

Z_k = ssp.coo_matrix( ( np.ones(p), ( np.where(B_k>0)[0], range(p) ) ), 
    shape=(mn, p), dtype=np.int8)
Z_u = ssp.coo_matrix( ( np.ones(q), ( np.where(B_u>0)[0], range(q) ) ), 
    shape=(mn, q), dtype=np.int8)
x_k = Z_k.T@X.flatten('F')

plt.figure(figsize=(16,9), dpi=100)
plt.subplot(121)
plt.imshow((Z_k@x_k).reshape(m,n, order='F'), cmap='gray')
plt.axis('off')
plt.title('Known part')
plt.subplot(122)
plt.imshow((Z_u@np.ones(q)).reshape(m,n, order='F'), cmap='gray')
plt.axis('off')
plt.title('Unknown part')
plt.show()

Dx = ssp.hstack( [ np.zeros( (m*(n-1),m) ), ssp.eye( m*(n-1) ) ] ) \
   - ssp.hstack( [ ssp.eye( m*(n-1) ), np.zeros( (m*(n-1),m) ) ] )
Dy = ssp.kron( ssp.eye(n), np.diff(np.eye(m), axis=0) )

DD = ssp.vstack( [ Dx, Dy ] )

import scipy.sparse.linalg as sla

A_MATRIX = DD@Z_u
B_MATRIX = DD@Z_k@x_k

x_u = sla.lsqr(A_MATRIX, -B_MATRIX)[0]
X_recon_L2 = (Z_k@x_k + Z_u@x_u).reshape(m, n, order='F')

plt.figure(figsize=(16,9), dpi=100)
plt.subplot(121)
plt.imshow(X_recon_L2, cmap='gray')
plt.title('Reconstructed image (L2)')
plt.axis('off')
plt.show()

import cvxpy as cp

v = cp.Variable(q)
obj = cp.Minimize( cp.norm( A_MATRIX@v + B_MATRIX, 1 ) )
prob = cp.Problem(obj)
prob.solve(verbose=True)
x_u = v.value

X_recon_L1 = (Z_k@x_k + Z_u@x_u).reshape(m, n, order='F')

plt.figure(figsize=(16,9), dpi=100)
plt.subplot(121)
plt.imshow(X_recon_L1, cmap='gray')
plt.title('Reconstructed image (L1)')
plt.axis('off')
plt.show()

X_clean = Image.open(requests.get("http://jonghank.github.io/ase7030/files/aerospace_building_clean.png", stream=True).raw)
X_clean = np.array(X_clean)/255

plt.figure(figsize=(16,9), dpi=100)
plt.subplot(221)
plt.imshow(X_clean, cmap='gray')
plt.title('Original image')
plt.axis('off')
plt.subplot(222)
plt.imshow(X_clean, cmap='gray')
plt.title('Original image')
plt.axis('off')
plt.show()

plt.figure(figsize=(16,9), dpi=100)
plt.subplot(221)
plt.imshow(X_recon_L2, cmap='gray')
plt.title('Reconstructed image (L2)')
plt.axis('off')
plt.subplot(222)
plt.imshow(X_recon_L1, cmap='gray')
plt.title('Reconstructed image (L1)')
plt.axis('off')
plt.subplot(223)
plt.imshow(np.abs(X_clean-X_recon_L2), cmap='gray')
plt.title('Difference image (L2 reconstruction)')
plt.axis('off')
plt.subplot(224)
plt.imshow(np.abs(X_clean-X_recon_L1), cmap='gray')
plt.title('Difference image (L1 reconstruction)')
plt.axis('off')
plt.show()