import enum
import pathlib

import numpy as np
from PIL import Image, ImageDraw


class UserInput(str, enum.Enum):
    SUM = '1'
    AVG = '2'
    MAX = '3'
    MIN = '4'
    CYCLE_MASK = '5'
    SQUARE_MASK = '6'
    RECTANGLE_MASK = '7'
    EXIT = '0'
    MULT = '__mult'


def _sum_tuple(size_1: tuple[int, int, int], size_2: tuple[int, int, int]) -> tuple[int, int, int]:
    return tuple(np.array(size_1) + np.array(size_2))


def _mult_tuple(size_1: tuple[int, int, int], size_2: tuple[int, int, int]) -> tuple[int, int, int]:
    return tuple(np.array(size_1) * np.array(size_2))


def _avg_tuple(size_1: tuple[int, int, int], size_2: tuple[int, int, int]) -> tuple:
    return tuple(map(int, (np.array(size_1) + np.array(size_2)) / 2))


def _max_tuple(size_1: tuple[int, int, int], size_2: tuple[int, int, int]) -> tuple[int, int, int]:
    return tuple(np.maximum.reduce([size_1, size_2]))


def _min_tuple(size_1: tuple[int, int, int], size_2: tuple[int, int, int]) -> tuple[int, int, int]:
    return tuple(np.minimum.reduce([size_1, size_2]))


def _cycle_mask(size: tuple[int, int]) -> Image.Image:
    mask: Image.Image = Image.new('RGBA', size, color=(0, 0, 0, 0))
    draw: ImageDraw.ImageDraw = ImageDraw.Draw(mask)

    width, height = size
    center = width // 2
    draw.ellipse(
        (
            center - width // 3,
            center - height // 3,
            center + width // 3,
            center + height // 3,
        ),
        fill=(1, 1, 1, 0),
    )
    return mask


def _square_mask(size: tuple[int, int]) -> Image.Image:
    mask: Image.Image = Image.new('RGBA', size, color=(0, 0, 0, 0))
    draw: ImageDraw.ImageDraw = ImageDraw.Draw(mask)

    width, height = size
    center = width // 2
    draw.rectangle(
        (
            center - width // 3,
            center - height // 3,
            center + width // 3,
            center + height // 3,
        ),
        fill=(1, 1, 1, 0),
    )
    return mask


def _rectangle_mask(size: tuple[int, int]) -> Image.Image:
    mask: Image.Image = Image.new('RGBA', size, color=(0, 0, 0, 0))
    draw: ImageDraw.ImageDraw = ImageDraw.Draw(mask)

    width, height = size
    center = width // 2
    draw.rectangle(
        (
            center - width // 2,
            center - height // 3,
            center + width // 2,
            center + height // 3,
        ),
        fill=(1, 1, 1, 0),
    )
    return mask


def _get_mask_by_user_input(user_input: UserInput, size: tuple[int, int]) -> Image.Image:
    mapper = {
        UserInput.RECTANGLE_MASK: _rectangle_mask,
        UserInput.SQUARE_MASK: _square_mask,
        UserInput.CYCLE_MASK: _cycle_mask,
    }
    mask_getter = mapper.get(user_input, UserInput.CYCLE_MASK)
    mask = mask_getter(size)
    return mask


def _process_command(
        user_input: UserInput,
        img_in_1: Image.Image,
        img_in_2: Image.Image,
        xy: tuple[int, int],
) -> tuple[int, int, int]:
    func_mapper = {
        UserInput.SUM: _sum_tuple,
        UserInput.AVG: _avg_tuple,
        UserInput.MAX: _max_tuple,
        UserInput.MIN: _min_tuple,
        UserInput.MULT: _mult_tuple,
    }

    img_1_pixel_slicer = slice(len(img_in_1.getpixel((0, 0))))
    img_2_pixel_slicer = slice(len(img_in_2.getpixel((0, 0))))
    if len(img_in_1.getpixel((0, 0))) > 3:
        img_1_pixel_slicer = slice(3)
    if len(img_in_2.getpixel((0, 0))) > 3:
        img_2_pixel_slicer = slice(3)

    return func_mapper[user_input](img_in_1.getpixel(xy)[img_1_pixel_slicer], img_in_2.getpixel(xy)[img_2_pixel_slicer])


def process_image(user_input: UserInput, img_in_1: Image.Image, img_in_2: Image.Image, out: Image.Image):
    width, height = img_in_1.size
    for x in range(width):
        for y in range(height):
            out.putpixel((x, y), _process_command(user_input, img_in_1, img_in_2, (x, y)))


def is_mask(user_input: UserInput) -> bool:
    return user_input in {UserInput.CYCLE_MASK, UserInput.SQUARE_MASK, UserInput.RECTANGLE_MASK}


def process_mask(user_input: UserInput, img_in_1: Image.Image, out: Image.Image):
    width, height = img_in_1.size
    mask = _get_mask_by_user_input(user_input, img_in_1.size)
    for x in range(width):
        for y in range(height):
            out.putpixel((x, y), _process_command(UserInput.MULT, img_in_1, mask, (x, y)))


def user_input_loop():
    while True:
        print(
            '1> Sum\n'
            '2> Avg\n'
            '3> Max\n'
            '4> Min\n'
            '5> Cycle mask\n'
            '6> Square mask\n'
            '7> Rectangle mask\n'
            '0> Exit\n'
        )

        try:
            user_input = UserInput(input('> '))
            return user_input
        except ValueError:
            print('Wrong option')
            continue


def get_init_images() -> tuple[Image.Image, Image.Image, Image.Image]:
    while True:
        image_1 = pathlib.Path(input('Enter first image path > '))
        if not image_1.exists():
            print('No such file.')
            continue
        break

    while True:
        image_2 = pathlib.Path(input('Enter second image path > '))
        if not image_2.exists():
            print('No such file.')
            continue

        break

    img_in_1: Image.Image = Image.open(image_1.resolve())
    img_in_2: Image.Image = Image.open(image_2.resolve())
    out: Image.Image = Image.new('RGBA', img_in_1.size, color=(0, 0, 0, 0))
    return img_in_1, img_in_2, out


def save_image_prompt(out: Image.Image):
    while True:
        path = pathlib.Path(input('Enter output image save path > '))
        if path.exists():
            print('This pass has been already taken.')
            continue
        try:
            out.save(path.open('wb'))
            return
        except OSError as e:
            print(e)
            path.unlink(missing_ok=True)


def main():
    im_1, im_2, out = get_init_images()
    user_input = user_input_loop()
    if user_input == UserInput.EXIT:
        return

    if is_mask(user_input):
        process_mask(user_input, im_1, out)
    else:
        if im_1.size != im_2.size:
            print('Image should be the same size!')
            return

        process_image(user_input, im_1, im_2, out)

    save_image_prompt(out)
    out.show()


if __name__ == '__main__':
    main()