import numpy as np


def is_green(i, j):
    return (not i%2 and not j%2) or (i%2 and j%2)

def is_red(i, j):
    return not i%2 and j%2

def is_blue(i, j):
    return i%2 and not j%2



def MDWI(image):
    
    res = np.empty((image.shape[0], image.shape[1], 3))
    estimated_green = np.zeros_like(image)
    estimated_red = np.zeros_like(image)
    estimated_blue = np.zeros_like(image)

    EPSILON = 10e-4

    # Estimation of green plan first

    pre_estimation = {}
    diag_gradient_factor = {}
    weight = {}
    H = [-8/256, 23/256, -48/256, 161/256, 161/256, -48/256, 23/256, -8/256]
    coord_NE = [(-4, -3), (-3, -2), (-2, -1), (-1, 0), (0, 1), (1, 2), (2, 3), (3, 4)]
    coord_SE = [(-3, -4), (-2, -3), (-1, -2), (0, -1), (1, 0), (2, 1), (3, 2), (4, 3)]
    coord_NW = [(3, -4), (2, -3), (1, -2), (0, -1), (-1, 0), (-2, 1), (-3, 2), (-4, 3)]
    coord_SW = [(4, -3), (3, -2), (2, -1), (1, 0), (0, 1), (-1, 2), (-2, 3), (-3, 4)]

    for i in range(3, image.shape[0] - 4):
        for j in range(3, image.shape[1] - 4):
            

            if (not is_green(i, j)):       
                pre_estimation["N"] = image[i-1, j] + (image[i, j] - image[i-2, j]) / 2   
                pre_estimation["S"] = image[i+1, j] + (image[i, j] - image[i+2, j]) / 2
                pre_estimation["W"] = image[i, j-1] + (image[i, j] - image[i, j-2]) / 2
                pre_estimation["E"] = image[i, j+1] + (image[i, j] - image[i, j+2]) / 2

                diag_gradient_factor["N"] = abs(image[i-2, j-1] - image[i, j-1]) + abs(image[i-3, j] - image[i-1, j]) + abs(image[i-2, j+1] - image[i, j+1]) + abs(image[i-3, j-1] - image[i-1, j-1]) + abs(image[i-3, j+1] - image[i-1, j+1]) + abs(image[i-2, j] - image[i, j]) + EPSILON
                diag_gradient_factor["S"] = abs(image[i+2, j-1] - image[i, j-1]) + abs(image[i+3, j] - image[i+1, j]) + abs(image[i+2, j+1] - image[i, j+1]) + abs(image[i+3, j-1] - image[i+1, j-1]) + abs(image[i+3, j+1] - image[i+1, j+1]) + abs(image[i+2, j] - image[i, j]) + EPSILON
                diag_gradient_factor["W"] = abs(image[i-1, j-2] - image[i-1, j]) + abs(image[i, j-3] - image[i, j-1]) + abs(image[i+1, j-2] - image[i+1, j]) + abs(image[i-1, j-3] - image[i-1, j-1]) + abs(image[i+1, j-3] - image[i+1, j-1]) + abs(image[i, j-2] - image[i, j]) + EPSILON
                diag_gradient_factor["E"] = abs(image[i-1, j+2] - image[i-1, j]) + abs(image[i, j+3] - image[i, j+1]) + abs(image[i+1, j+2] - image[i+1, j]) + abs(image[i-1, j+3] - image[i-1, j+1]) + abs(image[i+1, j+3] - image[i+1, j+1]) + abs(image[i, j+2] - image[i, j]) + EPSILON


                pre_estimation["NW"] = 0
                pre_estimation["SW"] = 0
                pre_estimation["NE"] = 0
                pre_estimation["SE"] = 0
                for k in range(8):
                    pre_estimation["NW"] += image[i + coord_NW[k][0], j + coord_NW[k][1]] * H[k]
                    pre_estimation["SW"] += image[i + coord_SW[k][0], j + coord_SW[k][1]] * H[k]
                    pre_estimation["NE"] += image[i + coord_NE[k][0], j + coord_NE[k][1]] * H[k]
                    pre_estimation["SE"] += image[i + coord_SE[k][0], j + coord_SE[k][1]] * H[k]
                
                diag_gradient_factor["NW"] = abs(image[i-2, j-1] - image[i-1, j]) + abs(image[i-1, j] - image[i, j+1]) + abs(image[i-1, j-2] - image[i, j-1]) + abs(image[i, j-1] - image[i+1, j]) + abs(image[i-1, j-1] - image[i+1, j+1]) + abs(image[i-2, j-2] - image[i, j]) + EPSILON
                diag_gradient_factor["SW"] = abs(image[i-1, j] - image[i, j-1]) + abs(image[i, j+1] - image[i+1, j]) + abs(image[i-1, j+1] - image[i+1, j-1]) + abs(image[i, j-1] - image[i+1, j-2]) + abs(image[i+1, j] - image[i+2, j-1]) + abs(image[i, j] - image[i+2, j-2]) + EPSILON
                diag_gradient_factor["NE"] = abs(image[i-1, j] - image[i, j-1]) + abs(image[i, j+1] - image[i+1, j]) + abs(image[i-1, j+1] - image[i+1, j-1]) + abs(image[i-1, j] - image[i-2, j+1]) + abs(image[i, j+1] - image[i-1, j+2]) + abs(image[i, j] - image[i-2, j+2]) + EPSILON
                diag_gradient_factor["SE"] = abs(image[i+2, j+1] - image[i+1, j]) + abs(image[i-1, j] - image[i, j+1]) + abs(image[i+1, j+2] - image[i, j+1]) + abs(image[i, j-1] - image[i+1, j]) + abs(image[i-1, j-1] - image[i+1, j+1]) + abs(image[i+2, j+2] - image[i, j]) + EPSILON

                
                for k, v in diag_gradient_factor.items():
                    weight[k] =  1 / v

                num = 0
                den = 0
                for k, v in pre_estimation.items():
                    num += v * weight[k]
                    den += weight[k]

                estimated_green[i, j] = num / den

            else:
                estimated_green[i, j] = image[i, j]


    # Then we work on blue and red plans
    
    diag = [(-1, -1), (-1, 1), (1, -1), (1, 1)]
    
    for i in range(3, image.shape[0] - 4):
        for j in range(3, image.shape[1] - 4):
            if (is_red(i, j) or is_blue(i, j)):
                ksi_nw = image[i-1, j-1] - estimated_green[i-1, j-1]
                ksi_ne = image[i-1, j+1] - estimated_green[i-1, j+1]
                ksi_sw = image[i+1, j-1] - estimated_green[i+1, j-1]
                ksi_se = image[i+1, j+1] - estimated_green[i+1, j+1]
                
                g_nw = abs(estimated_green[i-2, j-1] - estimated_green[i-1, j]) + abs(estimated_green[i-1, j-2] - estimated_green[i, j-1]) + abs(estimated_green[i-2, j-2] - estimated_green[i-1, j-1]) + abs(estimated_green[i-1, j-1] - estimated_green[i, j]) + abs(image[i-1, j-1] - image[i+1, j+1]) + EPSILON
                g_ne = abs(estimated_green[i-2, j+1] - estimated_green[i-1, j]) + abs(estimated_green[i-1, j+2] - estimated_green[i, j+1]) + abs(estimated_green[i-2, j+2] - estimated_green[i-1, j+1]) + abs(estimated_green[i-1, j+1] - estimated_green[i, j]) + abs(image[i-1, j+1] - image[i+1, j-1]) + EPSILON
                g_sw = abs(estimated_green[i+2, j-1] - estimated_green[i+1, j]) + abs(estimated_green[i+1, j-2] - estimated_green[i, j-1]) + abs(estimated_green[i+2, j-2] - estimated_green[i+1, j-1]) + abs(estimated_green[i+1, j-1] - estimated_green[i, j]) + abs(image[i+1, j-1] - image[i-1, j+1]) + EPSILON
                g_se = abs(estimated_green[i+2, j+1] - estimated_green[i+1, j]) + abs(estimated_green[i+1, j+2] - estimated_green[i, j+1]) + abs(estimated_green[i+2, j+2] - estimated_green[i+1, j+1]) + abs(estimated_green[i+1, j+1] - estimated_green[i, j]) + abs(image[i+1, j+1] - image[i-1, j-1]) + EPSILON

                pre_est = estimated_green[i, j] + (ksi_nw/g_nw + ksi_ne/g_ne + ksi_sw/g_sw + ksi_se/g_se) / (1/g_nw + 1/g_ne + 1/g_sw + 1/g_se)

                sum_dif = 0
                sum_sig = 0

                for coord in diag:
                    mu = image[i + coord[0], j + coord[1]] - pre_est
                    sum_dif += mu / (1 + abs(mu))
                    sum_sig += 1 / (1 + abs(mu))

                if (is_red(i, j)):
                    estimated_blue[i, j] = pre_est
                    estimated_red[i, j] = image[i, j]
                else:
                    estimated_red[i, j] = pre_est
                    estimated_blue[i, j] = image[i, j]

            else:
                if (is_red(i-1, j)):
                    estimated_red[i, j] = (image[i-1, j] + image[i+1, j]) / 2
                    estimated_blue[i, j] = (image[i, j-1] + image[i, j+1]) / 2
                else:
                    estimated_red[i, j] = (image[i, j-1] + image[i, j+1]) / 2
                    estimated_blue[i, j] = (image[i-1, j] + image[i+1, j]) / 2


    res[:, :, 0] = estimated_red
    res[:, :, 1] = estimated_green
    res[:, :, 2] = estimated_blue

    return res