import matplotlib.pyplot as plt
import numpy as np

    

def energy(img):
    (x,y,c) = img.shape
    r = np.zeros((x,y))
    for i in range(x):
        for j in range(y):
            r[i][j] = sum((img[(i-1) % x][j]-img[(i+1) % x][j])**2+(img[i][(j-1)%y]-img[i][(j+1)%y])**2)
    return r

def energy_vec(img):
    vertical = img[:-2] - img[2:]
    horizontal = img[:, :-2] - img[:, 2:]

    r = np.zeros(img.shape[:2])

    r[1:-1] += np.sum(vertical**2, axis=2)
    r[:, 1:-1] += np.sum(horizontal**2, axis=2)

    return r

def path_vec(energy):
    (h,w) = energy.shape
    paths = energy 
    
    # Calcul des chemins de moins énergies en tout point
    for i in range(h):
        paths[i] += np.min((paths[i-1], np.roll(paths[i-1], 1), np.roll(paths[i-1], -1)), axis = 0)
        paths[i][0] = energy[i][0] + min(paths[i-1][0], paths[i-1][1])
        paths[i][-1] = energy[i][-1] + min(paths[i-1][-1], paths[i-1][-2])

    # Reconstruction du chemin
    path = np.zeros(h, dtype=int)
    path[-1] = np.argmin(paths[-1])
    for i in range(1,h):
        j = path[h-i]
        if j == 0:
            path[h-i-1] = j + np.argmin((paths[h-i-1][0], paths[h-i-1][1]))
        elif j == w-1:
            path[h-i-1] = j + np.argmin((paths[h-i-1][-2], paths[h-i-1][-1])) - 1
        else:
            path[h-i-1] = j + np.argmin((paths[h-i-1][j-1], paths[h-i-1][j], paths[h-i-1][j+1])) - 1
    return path

def remove_path(img, p):
    h,w,b = img.shape
    r = np.zeros((h,w-1,b))
    for i in range(h):
        r[i] = np.delete(img[i], p[i], 0)
    return r

# For the culture: here is a vectorized version
def remove_path_vec(img, p):
    h, w, b = img.shape
    mask = np.ones((h, w), dtype=bool)
    mask[np.arange(h), p] = False
    return img[mask].reshape(h, w - 1, b)

def reduce_img(img, w): # On enlève w chemins dans img
    r = img
    for i in range(w):
        e=energy_vec(r)
        p=path_vec(e)
        r = remove_path_vec(r, p)
    return r

img = plt.imread("chateau.jpg") # img est non modifiable
img = img.astype(int) # transforme l'image en tableau modifiable de int

print(f"Dimensions originales: {img.shape}")
img = reduce_img(img, 200)
print(f"Dimensions finales: {img.shape}")

plt.imshow(img.astype(int)) # affichage de l'image dans une fenêtre
plt.show()