import torch


def discrete_haar_wavelet_2d(x, mode, N, levels=1):
    """ 2-D discrete Haar wavelet transform  

    Boundary conditions are resolved with periodic padding.

    The filters are normalized so that the forward and inverse mappings are
    unitary.
    """
    assert N % 2 == 0
    c = 0.5
 
    X = x.reshape((N,N))
    Y = torch.zeros_like(X)
    
    if mode == 1: # decomposition

        Np = N

        for level in range(levels):
            Nt = Np
            Np = Nt//2
            
            # approximation coeffs
            ll = c*(X[::2,::2]+X[::2,1::2]+X[1::2,::2]+X[1::2,1::2])
            # horizontal details
            lh = c*(X[::2,::2]+X[::2,1::2]-X[1::2,::2]-X[1::2,1::2])
            # vertical details
            hl = c*(X[::2,::2]-X[::2,1::2]+X[1::2,::2]-X[1::2,1::2])
            # diagonal details
            hh = c*(X[::2,::2]-X[::2,1::2]-X[1::2,::2]+X[1::2,1::2])
            
            Y[Np:Nt,Np:Nt] = hh
            Y[Np:Nt,:Np] = hl
            Y[:Np,Np:Nt] = lh
            X = ll

        Y[:Np,:Np] = X

        return Y.reshape(-1)

    else: # reconstruction

        Np = N//(2**levels)
        Nt = 2*Np

        for level in range(levels):
            ll = X[:Np,:Np]
            lh = X[:Np,Np:Nt]
            hl = X[Np:Nt,:Np]
            hh = X[Np:Nt,Np:Nt]

            Y[:Nt:2,:Nt:2] = c*(ll + lh + hl + hh)
            Y[:Nt:2,1:Nt:2] = c*(ll + lh - hl - hh)
            Y[1:Nt:2,:Nt:2] = c*(ll - lh + hl - hh)
            Y[1:Nt:2,1:Nt:2] = c*(ll - lh - hl + hh)
            
            X[:Nt,:Nt] = Y[:Nt,:Nt]

            Np = Nt
            Nt = 2*Np

        return Y.reshape(-1)