import numpy as np
import matplotlib.pyplot as plt
from src.checks import check_cfa, check_rgb 
from src.forward_model import CFA
import src.methods.Samanos_Thomas.version1 as v1
import src.methods.Samanos_Thomas.version2 as v2
from scipy import ndimage


def cfa_reconstruction(y: np.ndarray, cfa:str, input_shape: tuple) -> np.ndarray:
    """Performs demosaicking on y.

    Args:
        cfa (str): Name of the CFA. Can be bayer or quad_bayer.

    Returns:
        np.ndarray: Demosaicked image.
    """
    if cfa == "bayer":
        op = v2.bayer_reconstruction(y,input_shape)
    elif cfa == "quad_bayer":
        op = v2.quad_bayer_reconstruction(y,input_shape)
    else:
        op = np.zeros(input_shape)
    return op

def cfa_reconstruction_1(y: np.ndarray, cfa:str, input_shape: tuple) -> np.ndarray:
    """Performs demosaicking on y.

    Args:
        cfa (str): Name of the CFA. Can be bayer or quad_bayer.

    Returns:
        np.ndarray: Demosaicked image.
    """
    if cfa == "bayer":
        op = v1.bayer_reconstruction(y,input_shape)
    elif cfa == "quad_bayer":
        op = v1.quad_bayer_reconstruction(y,input_shape)
    else:
        op = np.zeros(input_shape)
    return op