import numpy as np
import itertools


def get_borders(image1, image2, shift_x, shift_y):
    if shift_x >= 0:
        x_left_first = shift_x
        x_left_second = 0
    else:
        x_left_first = 0
        x_left_second = -shift_x

    x_length = min(image1.shape[1] - x_left_first, image2.shape[1] - x_left_second)

    if shift_y >= 0:
        y_up_first = shift_y
        y_up_second = 0
    else:
        y_up_first = 0
        y_up_second = -shift_y

    y_length = min(image1.shape[0] - y_up_first, image2.shape[0] - y_up_second)

    return x_left_first, x_left_second, y_up_first, y_up_second, x_length, y_length



def mse(image1, image2, shift_x, shift_y):
    (x_left_first, x_left_second,
     y_up_first, y_up_second, x_length, y_length) = get_borders(image1, image2, shift_x, shift_y)

    subimage1 = image1[y_up_first: y_up_first + y_length,
                       x_left_first: x_left_first + x_length]
    subimage2 = image2[y_up_second: y_up_second + y_length,
                       x_left_second: x_left_second + x_length]
    return np.mean((subimage1 - subimage2) ** 2)


def find_optimal_shift(image1, image2, shift_variants_x, shift_variants_y, criterion):
    return min(itertools.product(shift_variants_x, shift_variants_y),
               key=lambda x: criterion(image1, image2, *x))


def pyramidal_alignment(image1, image2, criterion=mse):
    image1 = image1.astype(float)
    image2 = image2.astype(float)
    images = [(image1, image2)]
    while max(*image1.shape, *image2.shape) > 500:
        image1 = image1[::2, ::2]
        image2 = image2[::2, ::2]
        images.append((image1, image2))
    shift_x, shift_y = find_optimal_shift(image1, image2, range(-15, 16),
                                          range(-15, 16), criterion)
    images.pop()
    while images:
        image1, image2 = images.pop()
        shift_x *= 2
        shift_y *= 2
        shift_x, shift_y = find_optimal_shift(
            image1, image2,
            range(shift_x - 15, shift_x + 16),
            range(shift_y - 15, shift_y + 16),
            criterion,
        )
    return shift_x, shift_y


def cut(image):
    x_border = int(image.shape[1] * 0.05)
    y_border = int(image.shape[0] * 0.05)
    return image[y_border: -y_border, x_border: -x_border]


def align(image, green1):
    blue = cut(image[:len(image) // 3])
    green = cut(image[len(image) // 3: len(image) // 3 * 2])
    red = cut(image[len(image) // 3 * 2:])
    blue_x, blue_y = pyramidal_alignment(green, blue)
    red_x, red_y = pyramidal_alignment(green, red)

    (green_x1, blue_x, green_y1,
     blue_y, x_length_blue, y_length_blue) = get_borders(green, blue, blue_x, blue_y)
    (green_x2, red_x, green_y2,
     red_y, x_length_red, y_length_red) = get_borders(green, red, blue_x, blue_y)
    if green_x1 >= green_x2:
        green_x = green_x1
        red_x += green_x1 - green_x2
        x_length_red -= green_x1 - green_x2
    else:
        green_x = green_x2
        blue_x += green_x2 - green_x1
        x_length_blue -= green_x2 - green_x1

    if green_y1 >= green_y2:
        green_y = green_y1
        red_y += green_y1 - green_y2
        y_length_red -= green_y1 - green_y2
    else:
        green_y = green_y2
        blue_y += green_y2 - green_y1
        y_length_blue -= green_y2 - green_y1

    x_length = min(x_length_blue, x_length_red)
    y_length = min(y_length_blue, y_length_red)

    return (
        np.stack((red[red_y: red_y + y_length, red_x: red_x + x_length],
                  green[green_y: green_y + y_length, green_x: green_x + x_length],
                  blue[blue_y: blue_y + y_length, blue_x: blue_x + x_length]), axis=-1),
        (green1[0] + blue_y, green1[1] + blue_x), (green1[0] + red_y, green1[1] + red_x),
    )