'''
Produces low rank approximations of an image
The approximations will open as separate windws (see taskbar)
Also displays a plot of how the singular values decrease
Press any key when an image windows is active to quit the scipt
'''

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


#Image file, can replace by any image in the same folder as this script
file = "macaw.jpg"

#Generate pictures of the following ranks, needs to not exceed image dimensions
rank_list = [5, 20, 100]


image = cv2.imread(file, cv2.IMREAD_GRAYSCALE)

# Get the dimensions of the image
height, width = image.shape

#If the image is large you may want to resize it here:
new_dimensions = (width, height) 
print("New image dimensions: " , new_dimensions)

# Resize the image (if necessary)
small_image = cv2.resize(image, new_dimensions, interpolation=cv2.INTER_AREA)

def low_rank_approximation(image, rank=10):
    # Perform SVD
    U, S, Vt = np.linalg.svd(image, full_matrices=False)

    # Keep only the first 'rank' singular values and vectors
    U = U[:, :rank]
    S = np.diag(S[:rank])
    Vt = Vt[:rank, :]

    # Reconstruct the image
    approx_image = U @ S @ Vt

    # Ensure the pixel values are within the valid range
    approx_image[approx_image < 0] = 0
    approx_image[approx_image > 255] = 255
    approx_image = approx_image.astype('uint8')

    return approx_image

# Plotting singular values
_, singular_values, _ = np.linalg.svd(small_image, full_matrices=False)
plt.figure()
plt.plot(singular_values, 'o')  # 'o' for points
plt.title('Decay of Singular Values')
plt.xlabel('Index')
plt.ylabel('Singular Value')
plt.grid(True)
plt.show()

cv2.imshow('original', small_image)
cv2.imwrite("maccaw/maccaw_original.jpg", small_image)


for r in rank_list:
    im = low_rank_approximation(small_image, r)
    save_path = f"maccaw/maccaw_{r}.jpg"
    cv2.imwrite(save_path, im)
    cv2.imshow(f'Rank {r}', im)

cv2.waitKey(0)
cv2.destroyAllWindows()