"""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.laporte_auriane.fct_to_demosaick import *


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.
    """
    
    #print(f'y shape: {y.shape}') # (1024, 1024)
    input_shape = (y.shape[0], y.shape[1], 3)
    op = CFA(cfa, input_shape)
    
    ########## TODO ##########
    
     # Colour channels
    red = 0
    green = 1
    blue = 2
    correction = 0
    
    # In case of Quad_bayer op
    if op.cfa == 'quad_bayer':
        print("Transformation into Bayer pattern")
        y = quad_to_bayer(y, op)
        print('Quad_bayer to Bayer done\n\nDemoisaicking processing...')
 
    z = op.adjoint(y)
    # 1st and 2nd dim of a colour channel
    X = z[:,:,0].shape[0]
    Y = z[:,:,0].shape[1]
    #print(f"X: {X}") # 1024
    #print(f"Y: {Y}") # 1024
    
    
    ### Kimmel algorithm ###
    # About 3 to 4 minutes
    
    
    # Interpolation of green colour
    for i in range (X):
        for j in range (Y):
            # Control if value 0 in green channel
            if z[i,j,1] == 0:

                # Find neigbours
                n = neigh(i, j, green, z)
                # Derivatives
                d = derivat(n)
                # Weighting fct
                w = weight_fct(i, j, n, d, green, z)
                # Green interpolation (cross)
                z[i,j,1] = interpol_green(n, w)
    # Clip (avoid values our of range 0-1)
    z = clip_extreme(z)
    
    
    # 2 steps for the blues
    # Interpolate missing blues at the red locations
    for i in range (0, X, 2):
        for j in range (1, Y, 2):
            
            # Find neigbours + green
            n = neigh(i, j, blue, z)
            n_G = neigh(i, j, green, z)
            d = derivat(n_G)
            w = weight_fct(i, j, n, d, green, z)
            # Blue interpolation (4 corners)
            z[i,j,2] = interpol_red_blue(n, w, n_G)
    # Interpolate missing blues at the green locations
    for i in range (X):
        for j in range (Y):
            # Control if value 0 in blue channel
            if z[i,j,2] == 0:

                # Find neigbours
                n_B = neigh(i, j, blue, z)
                d = derivat(n_B)
                w = weight_fct(i, j, n_B, d, blue, z)
                # Blue interpolation (cross)
                z[i,j,2] = interpol_green(n_B,w)
    z = clip_extreme(z)


    # 2 steps for the reds
    # Interpolate missing reds at the blue locations
    for i in range (1, X, 2):
        for j in range (0, Y, 2):
            
            # Find neigbours + green
            n = neigh(i, j, red, z)
            n_G = neigh(i, j, green, z)
            d = derivat(n_G)
            w = weight_fct(i, j, n_G, d, green, z)
            # Red interpolation (4 corners)
            z[i,j,0] = interpol_red_blue(n, w, n_G)
    # Interpolate missing reds at the green locations
    for i in range (X):
        for j in range (Y):
            # Control if value 0 in red channel
            if z[i,j,0] == 0:

                # Find neigbours
                n_R = neigh(i, j, red, z)
                d = derivat(n_R)
                w = weight_fct(i, j, n_R, d, red, z)
                # Red interpolation (cross)
                z[i,j,0] = interpol_green(n_R,w)
    z = clip_extreme(z)
    
    
    # Correction step (repeated 3 times)
    if correction == 1:
        for i in range(3):
            
            n_G[4] = green_correction(n_R, n_G, n_B, w)
            n_B[4] = blue_correction(n_G, n_B, w)
            n_R[4] = red_correction(n_R, n_G, w)
            # nan
    
    #print(f'Image reconstructed shape: {z.shape}') # 1024, 1024, 3

    return z


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

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

# 2023
# Authors: Mauro Dalla Mura and Matthieu Muller