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

def bilinear_demosaicing(op: CFA, y: np.ndarray) -> np.ndarray:
    """
    Bilinear demosaicing method.

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

    Returns:
        np.ndarray: Demosaicked image.
    """
    # Copie des valeurs directement connues pour chaque canal
    red = y[:, :, 0]
    green = y[:, :, 1]
    blue = y[:, :, 2]

    # Création des masques pour chaque couleur selon le motif CFA
    mask_red = (op.mask == 0)  # Supposons que 0 correspond au rouge dans le masque
    mask_green = (op.mask == 1)  # Supposons que 1 correspond au vert
    mask_blue = (op.mask == 2)  # Supposons que 2 correspond au bleu

    # Interpolation bilinéaire pour le rouge et le bleu
    # Note: np.multiply multiplie les éléments correspondants des tableaux, c'est pourquoi nous utilisons np.multiply au lieu de *
    red_interp = convolve2d(np.multiply(red, mask_red), [[1/4, 1/2, 1/4], [1/2, 1, 1/2], [1/4, 1/2, 1/4]], mode='same')
    blue_interp = convolve2d(np.multiply(blue, mask_blue), [[1/4, 1/2, 1/4], [1/2, 1, 1/2], [1/4, 1/2, 1/4]], mode='same')

    # Interpolation bilinéaire pour le vert
    # Pour le vert, nous utilisons un autre noyau car il y a plus de pixels verts
    green_interp = convolve2d(np.multiply(green, mask_green), [[0, 1/4, 0], [1/4, 1, 1/4], [0, 1/4, 0]], mode='same')

    # Création de l'image interpolée
    demosaicked_image = np.stack((red_interp, green_interp, blue_interp), axis=-1)

    # Correction des valeurs interpolées: on réapplique les valeurs connues pour éviter le flou
    demosaicked_image[:, :, 0][mask_red] = red[mask_red]
    demosaicked_image[:, :, 1][mask_green] = green[mask_green]
    demosaicked_image[:, :, 2][mask_blue] = blue[mask_blue]

    # Clip pour s'assurer que toutes les valeurs sont dans la plage [0, 1]
    demosaicked_image = np.clip(demosaicked_image, 0, 1)

    return demosaicked_image


def quad_bayer_demosaicing(op: CFA, y: np.ndarray) -> np.ndarray:
    """
    Demosaicing method for Quad Bayer CFA pattern.

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

    Returns:
        np.ndarray: Demosaicked image.
    """
    
    # Interpolation bilinéaire pour chaque canal
    red_interp = convolve2d(np.multiply(y[:, :, 0], op.mask == 0), [[1/4, 1/2, 1/4], [1/2, 1, 1/2], [1/4, 1/2, 1/4]], mode='same')
    green_interp = convolve2d(np.multiply(y[:, :, 1], op.mask == 1), [[0, 1/4, 0], [1/4, 1, 1/4], [0, 1/4, 0]], mode='same')
    blue_interp = convolve2d(np.multiply(y[:, :, 2], op.mask == 2), [[1/4, 1/2, 1/4], [1/2, 1, 1/2], [1/4, 1/2, 1/4]], mode='same')

    # Assemblage de l'image interpolée
    demosaicked_image = np.stack((red_interp, green_interp, blue_interp), axis=-1)

    # Réapplication des valeurs connues
    demosaicked_image[:, :, 0][op.mask == 0] = y[:, :, 0][op.mask == 0]
    demosaicked_image[:, :, 1][op.mask == 1] = y[:, :, 1][op.mask == 1]
    demosaicked_image[:, :, 2][op.mask == 2] = y[:, :, 2][op.mask == 2]

    # Clip des valeurs pour les maintenir dans la plage appropriée
    demosaicked_image = np.clip(demosaicked_image, 0, 1)

    return demosaicked_image