import numpy as np
from matplotlib.image import imread
import matplotlib.pyplot as plt

import torch
from torchvision import transforms
from torchvision.datasets import MNIST

# the data is not normalized, and needs to be converted to a np array...
digits = MNIST('data', train=True, download=False)

# converting to np array and normalizing by val/max(arr)
digits = digits.data.numpy()
digits = digits/np.max(digits)

#displaying image
display = plt.imshow(digits[0])
display.set_cmap('gray')
print(digits[0].shape)

# grab 4x4 squares of the image
im1 = digits[0]

stride = 1
mat_sq = np.zeros([28,28,4,4])

for i in range(0, len(im1)-4, stride):
    for j in range(0, len(im1[i])-4, stride):
        cube = im1[i:i+4, j:j+4]
        mat_sq[i, j] = cube
        
# creates a 2d collection of matrices, where each [i, j] index contains a 4x4 matrix 
print(mat_sq[14, 14])
display = plt.imshow(mat_sq[14, 14])
display.set_cmap('gray')

# find SVD of each submatrix, make unitary and assign to uni
uni = np.zeros([28,28,4,4])

for i in range(len(mat_sq)):
    for j in range(len(mat_sq[i])):
        U, S, VT = np.linalg.svd(mat_sq[i, j])
        uni[i, j] = np.dot(U, VT)
        #print(f"U =\n{U}\n S =\n{S}\n VT=\n{VT}\n")

test = uni[14,14]
# comparing unitary, they aren't the same :(
print(f"U*VT:\n{uni[14, 14]}\nOriginal:\n{mat_sq[14, 14]}")


# try graham schmidt, they're just about the same!
# seems to have detected an edge.
Q, R = np.linalg.qr(mat_sq[14, 14])

print(f"Q:\n{Q}\nOriginal:\n{mat_sq[14, 14]}\nIs Q Unitary:\n{np.matmul(Q, Q.transpose(0, 1)).round()}")

display = plt.imshow(Q)
display.set_cmap('gray')