import numpy as np

import matplotlib.pyplot as plt

plt.rc('font', size=16, family='serif')

from PIL import Image, ImageOps

import urllib.request
urllib.request.urlretrieve("https://raw.githubusercontent.com/juanjosegarciaripoll/uimp-2020/master/cat.png", "cat.png")

cat = ImageOps.grayscale(Image.open('cat.png'))

cat

cat_array = np.asarray(cat)

cat_array[1,1]

cat_array[100,300]

(np.min(cat_array), np.max(cat_array))

cat_I = cat_array.flatten()+0.0
cat_ψ = np.sqrt(cat_I)/np.sqrt(np.sum(cat_I))

np.linalg.norm(cat_ψ)

def order_by_size(ψ, q):
    size = [2 for i in range(2*q)]
    ndx = [i + j for i in range(q) for j in [0, q]]
    return ψ.reshape(size).transpose(ndx)

def order_by_dimension(ψ, q):
    size = [2 for i in range(2*q)]
    ndx = [2*i + j for j in [0,1] for i in range(q)]
    return ψ.reshape(size).transpose(ndx)

def image2state(image, reorder=True):
    
    # Convert image to an array
    image = ImageOps.grayscale(image)
    image_I = np.asarray(image)
    
    # Reshape to fit some qubit size, the same for both dimensions
    L1, L2 = image_I.shape
    q1 = int(round(np.log2(L1)))
    q2 = int(round(np.log2(L2)))
    q = max(q1, q2)
    L = 2**q
    new_I = np.zeros((L, L), dtype=np.float64)
    new_I[0:L1,0:L2] = image_I
    
    I_t = np.sum(new_I.flatten())
    ψ = np.sqrt(new_I.flatten())/np.sqrt(I_t)
    
    if reorder:
        ψ = order_by_size(ψ, q)
    
    return ψ

def intensity2image(I):
    I = np.uint8(np.round(I/np.max(I) * 255.0))
    return Image.fromarray(I)

def state2image(ψ, reorder=True):
    q = int(round(np.log2(ψ.size))) // 2
    L = 2**q
    if reorder:
        ψ = order_by_dimension(ψ, q)
    return intensity2image(np.abs(ψ*ψ).reshape((L,L)))

state2image(image2state(cat, reorder=True), reorder=True)

def resample(ψ, nqubits=None, reorder=True):
    q = int(np.round(np.log2(ψ.size)))
    if nqubits is None:
        nqubits = q-2
    elif nqubits < 0:
        nqubits = q + nqubits
    
    L1 = 2**nqubits
    L2 = ψ.size // L1
    
    #
    # We are now going to extract the diagonal of
    # the reduced density matrix, and output it as a
    # new image
    if True:
        I = np.einsum('kii->k',
                      ψ.reshape((L1, L2, 1)) * ψ.conj().reshape(L1, 1, L2))
    else:
        # This is very slow and consumes a lot of memory
        ρ = ψ.reshape((L1,L2)) @ ψ.reshape((L1,L2)).T.conj()
        I = np.diag(ρ)

    # We may need to undo the ordering of qubits
    I = order_by_dimension(I, nqubits//2)
        
    # Intensity is now a vector. We have to rearrange it into
    # a square matrix for converting it to an image
    L = 2**(nqubits//2)
    return intensity2image(I.reshape((L,L)))

resample(image2state(cat), -2)

resample(image2state(cat), -4)

resample(image2state(cat), -6)

resample(image2state(cat), -8)

def enlarge(ψ, nqb):
    
    ψ = ψ.reshape((ψ.size,1)) * np.ones((1,2**nqb))
    ψ = ψ.flatten()
    return ψ / np.linalg.norm(ψ)

state2image(enlarge(image2state(resample(image2state(cat), -8)), 8))

def Gaussian(nqb, σ=0.1, reorder=True):
    
    x = np.linspace(-1, 1, 2**nqb)
    f = np.exp(-0.5 * (x/σ)**2)
    f /= np.linalg.norm(f)
    
    ψ = f.reshape((1,2**nqb)) * f.reshape((2**nqb,1))
    
    if reorder:
        ψ = order_by_size(ψ, nqb)
    
    return ψ / np.linalg.norm(ψ.flatten())

aux = state2image(Gaussian(9, 0.2))

aux

resample(image2state(aux), -2)

resample(image2state(aux), -4)

resample(image2state(aux), -6)

resample(image2state(aux), -8)

state2image(enlarge(image2state(resample(image2state(aux), -8)), 8))

def Schmidt(ψ, qubits_left=1):
    Lleft = 2**qubits_left
    Lright = ψ.size // Lleft
    ψ = ψ / np.linalg.norm(ψ)
    A = ψ.reshape((Lleft, Lright))
    
    Φ, sqrtΛ, Ξ = np.linalg.svd(A)
    return (sqrtΛ)**2

ψGaus = Gaussian(9, 0.2)
ψcat = image2state(cat)

fig, ax = plt.subplots()
ax.plot(Schmidt(ψcat, 9), label='cat')
ax.plot(Schmidt(ψGaus, 9), label='Gaussian')
ax.set_ylabel('$\\lambda_i$')
ax.set_xlabel('#eigenvalue')
ax.set_yscale('log')
ax.set_ylim([1e-12,1.2])
ax.legend();