"""The main file for the reconstruction.
This file should NOT be modified except the body of the 'run_reconstruction' function.
Students can call their functions (declared in others files of src/methods/your_name).
"""


import numpy as np

from src.forward_model import CFA
from src.methods.quelletl.some_function import interpolation
import cv2
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
from scipy.signal import convolve2d


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

    Args:
        y (np.ndarray): Mosaicked image to be reconstructed.
        cfa (str): Name of the CFA. Can be bayer or quad_bayer.

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

    input_shape = (y.shape[0], y.shape[1], 3)
    op = CFA(cfa, input_shape)
    res = interpolation(y, op)
    
    return res

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

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

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

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

# 2023
# Authors: Mauro Dalla Mura and Matthieu Muller