import numpy as np
import PIL
from kymatio.numpy import Scattering2D
from kymatio.visuals import imshow

def load_rgb(img):
    img = np.array(PIL.Image.open(path).convert("RGB")).astype('float64')
    img /= np.abs(img).max(axis=(0, 1))
    return img

def group_by_scale(out, J, L):
    n_S1 = J * L  # number of first-order coeffs
    S1_all = out[1:n_S1 + 1]
    S1_scale = S1_all.reshape(L, -1, *out.shape[-3:], order='F').mean(axis=0)
    return S1_scale

def viz_by_scale(S1_scale, J, ticks=0, second_order=False):
    J = J - 1 if second_order else J
    for j in range(J):
        j_title = j + 1 if second_order else j
        imshow(S1_scale[j], title=f"scale={j_title}",
               w=.5, h=.5, abs=1, ticks=ticks)

def unpad(out, ref):
    assert out.ndim == 4
    rx, ry = ref.shape[:2]
    ox, oy = out.shape[1:3]
    dx, dy = abs(ox - rx) // 2, abs(oy - ry) // 2
    return out[:, dx:-dx, dy:-dy]

#%%
path = r"C:\Desktop\colors.png"
path = r"C:\Desktop\lena.jpg"
img = load_rgb(path)
#%%
imshow(img, w=.6, h=.6, title="%s x %s image" % img.shape[:2])
#%%
J, L = 4, 8  # largest scale; number of angles
S = Scattering2D(shape=img.shape[:2], L=L, J=J)

#%%
# take 3 wavelet transforms on each channel and average
# out = np.mean([S(im) for im in img.transpose(-1, 0, 1)], axis=0)
outs = [S(im) for im in img.transpose(-1, 0, 1)]
#%%
out = np.vstack([o[None] for o in outs]).transpose(1, 2, 3, 0)
print(img.shape)
print(out.shape)
#%%
S1_scale = group_by_scale(out, J, L)
S1_scale /= S1_scale.max(axis=(0, 1, 2))
S1_scale = unpad(S1_scale, ref=img)
viz_by_scale(S1_scale, J)