import numpy as np
import pywt
import matplotlib.pyplot as plt

def generate_haar_basis_image(size: int, i: int, j: int, level: int = 1) -> np.ndarray:
    """
    Generate a basis image for the Haar wavelet transform by setting a single coefficient, and using the inverse transform.

    Parameters:
        size (int): The size of the image (height and width).
        i (int): The row index of the coefficient to set.
        j (int): The column index of the coefficient to set.
        level (int): The level of decomposition for the wavelet transform.

    Returns:
        np.ndarray: The reconstructed basis image after inverse Haar transform.
    """
    # Initialize the Haar wavelet and coefficient matrix
    haar_wavelet = pywt.Wavelet('haar')
    coeff_structure = pywt.wavedec2(np.zeros((size, size)), haar_wavelet, level=level)
    coeffs, coeff_slices = pywt.coeffs_to_array(coeff_structure)

    # Set the specific coefficient to 1 (rest are zeros)
    coeffs[i, j] = 1.0

    # Perform inverse transform to get the basis image
    reconstructed_coeffs = pywt.array_to_coeffs(coeffs, coeff_slices, output_format='wavedec2')
    basis_image = pywt.waverec2(reconstructed_coeffs, haar_wavelet)
    return basis_image

def generate_haar_basis_images(size: int, level: int = 1) -> list:
    """
    Generate all basis images for a Haar wavelet transform.

    Parameters:
        size (int): The size of the image (height and width).
        level (int): The level of decomposition for the wavelet transform.

    Returns:
        list: A list of all generated basis images.
    """
    return [generate_haar_basis_image(size, i, j, level=level) for i in range(size) for j in range(size)]


def plot_basis_images(basis_images, size, title, save_path):
    plt.figure(figsize=(10, 10))
    for i, basis_image in enumerate(basis_images):
        plt.subplot(size, size, i + 1)
        plt.imshow(basis_image, vmin=-.5, vmax=.5, cmap='gray')
        plt.axis('off')
    plt.suptitle(title)
    plt.tight_layout()
    plt.show()


# Parameters for the experiment
size = 8  # Image size: 8x8 (must here be a power of two)

# Show single-level Haar Transform basis images
basis_images_single = generate_haar_basis_images(size, level=1)
plot_basis_images(basis_images_single, size, "Basis Images for Single-Level Haar Transform", "single_level_basis_images.png")

# Show multi-level Haar Transform basis images
max_levels = int(np.log2(size))
basis_images_multi = generate_haar_basis_images(size, level=max_levels)
plot_basis_images(basis_images_multi, size, f"Basis Images for Multi-Scale Haar Transform ({max_levels} Levels)", "multi_scale_basis_images.png")

# Function to compute inner product between two basis images
def inner_product(image1, image2):
    return np.sum(image1 * image2)

# Function to compute the L2 norm of an image (basis image)
def compute_l2_norm(image):
    return np.sqrt(np.sum(image**2))

# Select two basis images (for example, the first and the second)
ix_1 = 1
ix_2 = 5
basis_image_1 = basis_images_multi[ix_1]
basis_image_2 = basis_images_multi[ix_2]

# Calculate the inner product
inner_prod = inner_product(basis_image_1, basis_image_2)
print(f'Inner product between basis image {ix_1} and {ix_2}: {inner_prod}')