import numpy as np
from scipy import signal
from src.forward_model import CFA

def compute_gradient(channel, points):
    """ Compute the gradient for the channel based on the specified points. """
    gradient = np.zeros_like(channel)
    for point in points:
        gradient += np.roll(channel, shift=point, axis=(0, 1))
    gradient /= len(points)
    gradient -= channel
    return gradient

def gradient_correction_interpolation(op: CFA, z: np.ndarray,alpha=0.5, beta=0.5, gamma=0.5) -> np.ndarray:
    # defining bilinear interpolation filters for the different channels (bayer pattern case)
    bilinear_filter_red = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]]) / 4
    bilinear_filter_green = np.array([[0, 1, 0], [1, 4, 1], [0, 1, 0]]) / 4
    bilinear_filter_blue = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]]) / 4

    interpolated_img = np.empty(op.input_shape)
    
    # applying the convolution product with the bilinear filters
    interpolated_img[:, :, 0] = signal.convolve2d(z[:, :, 0], bilinear_filter_red,  boundary='symm',mode='same')
    interpolated_img[:, :, 1] = signal.convolve2d(z[:, :, 1], bilinear_filter_green,  boundary='symm',mode='same')
    interpolated_img[:, :, 2] = signal.convolve2d(z[:, :, 2], bilinear_filter_blue, boundary='symm', mode='same')

    """ Applying gradient correction to the bilinearly interpolated image. """


    R = interpolated_img[:, :, 0]
    G = interpolated_img[:, :, 1]
    B = interpolated_img[:, :, 2]

    # computing gradients for R and B channels
    points_R = [(0, -2), (0, 2), (-2, 0), (2, 0)] 
    points_B = [(-1, -1), (-1, 1), (1, -1), (1, 1), (0, 0)]  
    grad_R = compute_gradient(R, points_R)
    grad_B = compute_gradient(B, points_B)

    G_corrected = np.copy(G)
    G_corrected[0::2, 0::2] = G[0::2, 0::2] + alpha * grad_R[0::2, 0::2]
    
    G_corrected[1::2, 1::2] = G[1::2, 1::2] + gamma * grad_B[1::2, 1::2]
    
    # correction of R at G locations and B at G locations
    R_corrected = np.copy(R)
    B_corrected = np.copy(B)
    R_corrected[1::2, 0::2] = R[1::2, 0::2] + beta * grad_R[1::2, 0::2]  
    R_corrected[0::2, 1::2] = R[0::2, 1::2] + beta * grad_R[0::2, 1::2] 
    B_corrected[1::2, 0::2] = B[1::2, 0::2] + beta * grad_B[1::2, 0::2] 
    B_corrected[0::2, 1::2] = B[0::2, 1::2] + beta * grad_B[0::2, 1::2]  

   
    corrected_image = np.stack((R_corrected, G_corrected, B_corrected), axis=-1)

    return corrected_image

def swapping(op, y):
    """
    Convert both the CFA operator's mask and an image from Quad Bayer to Bayer by swapping method.

    Args:
        op: CFA operator.
        y (np.ndarray): Mosaicked image.

    Returns:
        tuple: A tuple containing the updated CFA operator and the converted image.
    """

    if op.cfa == 'quad_bayer':
        #  swapping 2 columns every 2 columns
        op.mask[:, 1::4], op.mask[:, 2::4] = op.mask[:, 2::4].copy(), op.mask[:, 1::4].copy()
        #  swapping 2 lines every 2 lines
        op.mask[1::4, :], op.mask[2::4, :] = op.mask[2::4, :].copy(), op.mask[1::4, :].copy()
        # swapping the diagonal pairs
        op.mask[1::4, 1::4], op.mask[2::4, 2::4] = op.mask[2::4, 2::4].copy(), op.mask[1::4, 1::4].copy()

    converted_image = y.copy()
    converted_image[:, 1::4], converted_image[:, 2::4] = converted_image[:, 2::4].copy(), converted_image[:, 1::4].copy()
    converted_image[1::4, :], converted_image[2::4, :] = converted_image[2::4, :].copy(), converted_image[1::4, :].copy()
    converted_image[1::4, 1::4], converted_image[2::4, 2::4] = converted_image[2::4, 2::4].copy(), converted_image[1::4, 1::4].copy()

    # updating to bayer pattern
    op.cfa = 'bayer'

    return op, converted_image