import numpy as np
from scipy.signal import convolve2d

from src.forward_model import CFA

def bilinear_interpolation(op: CFA, z: np.ndarray) -> np.ndarray:
    """Perform bilinear interpolation for demosaicing

    Args:
        op (CFA): CFA operator.
        z (np.ndarray): Adjoint image.

    Returns:
        np.ndarray: Interpolated image.
    """
    # Bi-linear interpolation

    ker_bayer_red_blue = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]]) / 4
    ker_bayer_green = np.array([[0, 1, 0], [1, 4, 1], [0, 1, 0]]) / 4

    res = np.empty(op.input_shape)

    res[:, :, 0] = convolve2d(z[:, :, 0], ker_bayer_red_blue, mode='same')
    res[:, :, 1] = convolve2d(z[:, :, 1], ker_bayer_green, mode='same')
    res[:, :, 2] = convolve2d(z[:, :, 2], ker_bayer_red_blue, mode='same')

    return res

def spectral_difference(op: CFA, z: np.ndarray, res: np.ndarray) -> np.ndarray:
    """Perform spectral difference method for demosaicing

    Args:
        op (CFA): CFA operator.
        z (np.ndarray): Adjoint image.
        res (np.ndarray): Interpolated image.

    Returns:
        np.ndarray: Demosaicked image.
    """

    ker_bayer_red_blue = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]]) / 4
    ker_bayer_green = np.array([[0, 1, 0], [1, 4, 1], [0, 1, 0]]) / 4

    # Computation of spectral differences

    delta_red_green_quad =  z[:, :, 0] - np.multiply(res[:, :, 1],op.mask[:,:,0])
    delta_red_blue_quad =  z[:, :, 0] - np.multiply(res[:, :, 2],op.mask[:,:,0])

    delta_blue_green_quad =  z[:, :, 2] - np.multiply(res[:, :, 1],op.mask[:,:,2])
    delta_blue_red_quad =  z[:, :, 2] - np.multiply(res[:, :, 0],op.mask[:,:,2])

    delta_green_red_quad =  z[:, :, 1] - np.multiply(res[:, :, 0],op.mask[:,:,1])
    delta_green_blue_quad =  z[:, :, 1] - np.multiply(res[:, :, 2],op.mask[:,:,1])

    # Estimation
    res_sd = np.empty(op.input_shape)

    res_sd[:,:,0] = res[:, :, 1] + convolve2d(delta_red_green_quad,ker_bayer_red_blue,mode='same') + res[:, :, 2] + convolve2d(delta_red_blue_quad,ker_bayer_red_blue,mode='same') 
    res_sd[:,:,2] = res[:, :, 1] + convolve2d(delta_blue_green_quad,ker_bayer_red_blue,mode='same') + res[:, :, 0] + convolve2d(delta_blue_red_quad,ker_bayer_red_blue,mode='same') 
    res_sd[:,:,1] = res[:, :, 0] + convolve2d(delta_green_red_quad,ker_bayer_green,mode='same') + res[:, :, 2] + convolve2d(delta_green_blue_quad,ker_bayer_green,mode='same') 

    return res_sd

def normalization(res_sd:np.ndarray)-> np.ndarray:
    """Perform a min-max normalization

    Args:
        res_sd (np.ndarray): Demosaicked image.

    Returns:
        np.ndarray: Normalized image.
    """

    res = (res_sd - np.min(res_sd)) / (np.max(res_sd) - np.min(res_sd))
    
    return res

def quad_bayer_to_bayer_pattern(op: CFA):
    """Quad to Bayer pattern conversion by swapping method

    Args:
        op (CFA): CFA operator.

    Returns:
        CFA: Bayer pettern.
    """
    if op.cfa == 'quad_bayer':
        for j in range(1, op.mask.shape[1], 4):
            op.mask[:,j], op.mask[:,j+1] = op.mask[:,j+1].copy(), op.mask[:,j].copy()

        for i in range(1, op.mask.shape[0], 4):
            op.mask[i, :], op.mask[i+1,:] = op.mask[i+1,:].copy(), op.mask[i,:].copy()
            
        for i in range(1, op.mask.shape[0], 4):
            for j in range(1, op.mask.shape[1], 4):
                op.mask[i,j], op.mask[i+1,j+1] = op.mask[i+1,j+1].copy(), op.mask[i,j].copy() 
    return 0

def quad_bayer_to_bayer_image(y:np.array):
    """Quad to Bayer conversion of a mosaicked image by swapping method

    Args:
        y (np.array): Mosaicked image.
    """
        
    for j in range(1, y.shape[1], 4):
        y[:,j], y[:,j+1] = y[:,j+1].copy(), y[:,j].copy()

    for i in range(1, y.shape[0], 4):
        y[i, :], y[i+1,:] = y[i+1,:].copy(), y[i,:].copy()
        
    for i in range(1, y.shape[0], 4):
        for j in range(1, y.shape[1], 4):
            y[i,j], y[i+1,j+1] = y[i+1,j+1].copy(), y[i,j].copy() 
            
    return 0