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

def malvar(y: np.ndarray, op: CFA) -> np.ndarray:
    """
    Malvar-He-Cutler demosaicing algorithm for Bayer pattern.
    """
    # Convert the mosaicked image to the initial estimated channels.
    z = op.adjoint(y)

    # Define convolution kernels
    kernel_G_at_RB = np.array([[0, 0, -1, 0, 0], [0, 0, 2, 0, 0], [-1, 2, 4, 2, -1], [0, 0, 2, 0, 0], [0, 0, -1, 0, 0]]) / 8
    kernel_RB_at_G = np.array([[0, 0, 0.5, 0, 0], [0, -1, 0, -1, 0], [-1, 4, 5, 4, -1], [0, -1, 0, -1, 0], [0, 0, 0.5, 0, 0]]) / 8
    kernel_RB_at_RB = np.array([[0, 0, -1.5, 0, 0], [0, 2, 0, 2, 0], [-1.5, 0, 6, 0, -1.5], [0, 2, 0, 2, 0], [0, 0, -1.5, 0, 0]]) / 8

    # Interpolate each channel
    R = z[:, :, 0]
    G = z[:, :, 1]
    B = z[:, :, 2]

    mask = op.mask
    R_m, G_m, B_m = mask[:, :, 0], mask[:, :, 1], mask[:, :, 2]

    # Interpolate G at R and B locations
    G = np.where(np.logical_or(R_m == 1, B_m == 1), convolve2d(y, kernel_G_at_RB, mode='same'), G)

    # Interpolate R at G and B locations, B at R and G locations
    R = np.where(np.logical_or(G_m == 1, B_m == 1), convolve2d(y, kernel_RB_at_G, mode='same'), R)
    B = np.where(np.logical_or(R_m == 1, G_m == 1), convolve2d(y, kernel_RB_at_G, mode='same'), B)

    # Interpolate R at B locations and B at R locations
    R = np.where(B_m == 1, convolve2d(y, kernel_RB_at_RB, mode='same'), R)
    B = np.where(R_m == 1, convolve2d(y, kernel_RB_at_RB, mode='same'), B)

    # Combine channels
    return np.clip(np.stack((R, G, B), axis=-1), 0, 1)





####
####
####

####      ####                ####        #############
####      ######              ####      ##################
####      ########            ####      ####################
####      ##########          ####      ####        ########
####      ############        ####      ####            ####
####      ####  ########      ####      ####            ####
####      ####    ########    ####      ####            ####
####      ####      ########  ####      ####            ####
####      ####  ##    ######  ####      ####          ######
####      ####  ####      ##  ####      ####    ############
####      ####  ######        ####      ####    ##########
####      ####  ##########    ####      ####    ########
####      ####      ########  ####      ####
####      ####        ############      ####
####      ####          ##########      ####
####      ####            ########      ####
####      ####              ######      ####

# 2023
# Authors: Mauro Dalla Mura and Matthieu Muller