done ver1

from pickletools import uint8
import numpy as np
import cv2
import math
def show(image):
    cv2.imshow('image',image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

def get_gs(window_size, sigma_s, pad_w):
    '''return gs, size: (window_size, window_size)'''
    gs = np.zeros((window_size, window_size, 3))
    for i in range(window_size):
        for j in range(i, window_size):
            result =  np.exp((-1) * ((i-pad_w)**2 + (j-pad_w)**2) * 0.5  / (sigma_s**2))
            gs[i, j, :] = result 
            gs[j, i, :] = result
    return gs

# def get_gr_value(p, q, gr_dict):
#     p_red, p_green, p_blue = p[0], p[1], p[2]
#     q_red, q_green, q_blue = q[0], q[1], q[2]
#     gr = gr_dict[abs(p_red - q_red)] * gr_dict[abs(p_green -  q_green)] * gr_dict[abs(p_blue-q_blue)]
#     return gr



class Joint_bilateral_filter(object):
    def __init__(self, sigma_s, sigma_r):
        self.sigma_r = sigma_r
        self.sigma_s = sigma_s
        self.wndw_size = 6*sigma_s+1 # ws
        self.pad_w = 3*sigma_s # r
    
    def joint_bilateral_filter(self, img, guidance):
        BORDER_TYPE = cv2.BORDER_REFLECT
        padded_img = cv2.copyMakeBorder(img, self.pad_w, self.pad_w, self.pad_w, self.pad_w, BORDER_TYPE).astype(np.int32)
        padded_guidance = cv2.copyMakeBorder(guidance, self.pad_w, self.pad_w, self.pad_w, self.pad_w, BORDER_TYPE).astype(np.int32)
        window_size = self.wndw_size
        pad_w = self.pad_w
        print("window size: ", window_size)
        print("padded image size: ", padded_img.shape)
        print("padded guidance size: ", padded_guidance.shape)
        print("=============")
        
        ### TODO ###
        # save Gs kernel dict

        gs = get_gs(window_size, self.sigma_s, self.pad_w)
        gr_dict = [np.exp(-(x * 0.5 / ((self.sigma_r**2)* (255**2)))) for x in range(256)]
        

        output = np.zeros(img.shape)
        
        shift = 0
        i, j = 0, 0
        shift_x, shift_y = 0, 0
        
        # for a window
        for i in range(img.shape[0]): 
           for j in range(img.shape[1]):
                intensity = padded_img[i:i+window_size, j:j+window_size]
                Tq = padded_guidance[i:i+window_size, j:j+window_size]/255
                Tp = padded_guidance[i+pad_w, j+pad_w]/255
                if (len(Tq.shape) == 3):
                    # print(Tq-Tp)
                    # print((Tq-Tp).shape)
                    gr = np.exp(-np.sum(np.square(Tq-Tp), axis=2) * 0.5 / (self.sigma_r**2))
                else:
                    gr = np.exp(-np.square(Tq-Tp) * 0.5 / (self.sigma_r**2))
                gr = np.stack((gr, gr, gr), axis=-1)
                
                output[i, j] = np.sum(np.sum(gs*gr*intensity, axis=0), axis=0) / np.sum(np.sum((gs * gr), axis=0), axis=0)
                
        return np.clip(output, 0, 255).astype(np.uint8)