import numpy as np
# from scipy.fft import fft2, ifft2, fftshift, fftfreq, ifftshift
import imageio.v3 as imageio
import matplotlib.pyplot as plt
# import seaborn as sn
import os
import argparse

def luminance(l):
    if len(l.shape) > 2:
        l = 0.2126 * l[:, :, 0] + 0.7152 * l[:, :, 1] + 0.0722 * l[:, :, 2]
        l = l.astype(np.uint8)
    return l

def normalization(g):
    g = (((g - g.min()) * 255) / (g.max() - g.min()))
    return g.astype(np.uint8)

def bilateral_filter_patch(image, diameter, sigma_color, sigma_space):

    h, w = image.shape
    r = diameter // 2
    img_pad = np.pad(image, r, 'reflect')

    samples = np.meshgrid(np.arange(-r, +r+1), np.arange(-r, +r+1))
    gaussian_space = np.exp(-(samples[0]**2 + samples[1]**2)/(2*sigma_space**2))

    filtered_img = np.zeros_like(img_pad)

    filters = []
    regions = []
    for i in range(r, h+r):
        for j in range(r, w+r):
            region = img_pad[i-r:i+r+1, j-r:j+r+1]
            # breakpoint()
            center = img_pad[i, j]
            gaussian_int = np.exp(-((center - region)**2)/(2*sigma_color**2))
            filter_bilateral = gaussian_space*gaussian_int
            filtered_img[i, j] = np.sum(region*filter_bilateral)/np.sum(filter_bilateral)
            
            regions.append(region)
            filters.append(filter_bilateral)

    return filtered_img, filters, regions

def main():

    img = imageio.imread('ball.jpeg', pilmode='L').astype(np.float64)
    filtered, filters, regions = bilateral_filter_patch(img, diameter=55, sigma_color=12.0, sigma_space=16.0)

    fig, axes = plt.subplots(1, 2, figsize=(14, 7))
    ax = axes[0]
    ax.imshow(img, cmap='gray')
    ax = axes[1]
    ax.imshow(filtered, cmap='gray')
    plt.tight_layout()
    plt.show()

    num_rows = num_cols = int(np.sqrt(len(filters)))
    filters = np.reshape(np.array(filters), (num_rows, num_cols, 55, 55))
    regions = np.reshape(np.array(regions), (num_rows, num_cols, 55, 55))

    fig, axes = plt.subplots(3, 6, figsize=(14, 7))

    ax_filt = axes[0, 0]
    ax_filt.imshow(filters[100, 100], cmap='gray')
    ax_filt.axis('off')
    ax_filt = axes[0, 1]
    ax_filt.imshow(regions[100, 100], cmap='gray')
    ax_filt.axis('off')

    ax_filt = axes[0, 2]
    ax_filt.imshow(filters[200, 200], cmap='gray')
    ax_filt.axis('off')
    ax_filt = axes[0, 3]
    ax_filt.imshow(regions[200, 200], cmap='gray')
    ax_filt.axis('off')

    ax_filt = axes[0, 4]
    ax_filt.imshow(filters[300, 300], cmap='gray')
    ax_filt.axis('off')
    ax_filt = axes[0, 5]
    ax_filt.imshow(regions[300, 300], cmap='gray')
    ax_filt.axis('off')

    ax_filt = axes[1, 0]
    ax_filt.imshow(filters[150, 150], cmap='gray')
    ax_filt.axis('off')
    ax_filt = axes[1, 1]
    ax_filt.imshow(regions[150, 150], cmap='gray')
    ax_filt.axis('off')

    ax_filt = axes[1, 2]
    ax_filt.imshow(filters[250, 250], cmap='gray')
    ax_filt.axis('off')
    ax_filt = axes[1, 3]
    ax_filt.imshow(regions[250, 250], cmap='gray')
    ax_filt.axis('off')

    ax_filt = axes[1, 4]
    ax_filt.imshow(filters[20, 20], cmap='gray')
    ax_filt.axis('off')
    ax_filt = axes[1, 5]
    ax_filt.imshow(regions[20, 20], cmap='gray')
    ax_filt.axis('off')

    ax_filt = axes[2, 0]
    ax_filt.imshow(filters[110, 110], cmap='gray')
    ax_filt.axis('off')
    ax_filt = axes[2, 1]
    ax_filt.imshow(regions[110, 110], cmap='gray')
    ax_filt.axis('off')

    ax_filt = axes[2, 2]
    ax_filt.imshow(filters[220, 220], cmap='gray')
    ax_filt.axis('off')
    ax_filt = axes[2, 3]
    ax_filt.imshow(regions[220, 220], cmap='gray')
    ax_filt.axis('off')

    ax_filt = axes[2, 4]
    ax_filt.imshow(filters[50, 50], cmap='gray')
    ax_filt.axis('off')
    ax_filt = axes[2, 5]
    ax_filt.imshow(regions[50, 50], cmap='gray')
    ax_filt.axis('off')

    plt.tight_layout()
    plt.show()

if __name__ == '__main__':
    main()