import random
from io import BytesIO

import numpy as np
import multiprocessing

import requests
from mpi4py import MPI
from PIL import Image


def count_randomized_hits(iterations):
    counter = 0
    for i in range(iterations):
        x = random.uniform(0, 1)
        y = random.uniform(0, 1)
        if ((x ** 2) + (y ** 2)) ** (1 / 2) <= 1:
            counter += 1
    return counter


FLAG_default = 0
FLAG_threaded = 1
FLAG_network = 2


def monte_carlo_methode(n, mode=0):
    if mode == 1:  # Multithreading mode
        num_threads = 16
        iterations_per_thread = n // num_threads
        with multiprocessing.Pool(num_threads) as pool:
            hits = pool.map(count_randomized_hits, [iterations_per_thread] * num_threads)
            hit = sum(hits)

    elif mode == 2:  # MPI parallel mode
        comm = MPI.COMM_WORLD
        rank = comm.Get_rank()
        size = comm.Get_size()
        local_hit = count_randomized_hits(n // size)
        hit = comm.reduce(local_hit, op=MPI.SUM, root=0)

    else:  # Default mode
        hit = count_randomized_hits(n)

    pi_approx = (hit / n) * 4
    pi_diff = abs(np.pi - pi_approx)

    return pi_approx, pi_diff


def uniform_kernel(n):
    if n % 2 == 0:
        print("Size needs to be odd")
        exit(1)
    K = 1 / n / n * np.ones([n, n])
    return K


def gauss_kernel(s):
    n = 3 * s
    pos = np.arange(-n, n + 1)
    x = np.meshgrid(pos, pos)
    K = 1.0 / (2.0 * np.pi * s * s) * np.exp(-(x[0] ** 2 + x[1] ** 2) / (2.0 * s * s))
    K = K / sum(sum(K))
    return K


FLAG_gauss = 0
FLAG_uniform = 1


def process_image_part(data_part, kernel, padding):
    y_part_size, x_part_size, _ = data_part.shape

    data_part_new = np.zeros((data_part.shape[0] - padding[0], data_part.shape[1] - padding[1], 3))

    pad_y, pad_x = padding
    for i in range(pad_y, y_part_size - pad_y):
        for j in range(pad_x, x_part_size - pad_x):
            for k in range(3):
                new_value = 0.0
                for ii in range(kernel.shape[0]):
                    for jj in range(kernel.shape[1]):
                        iii = ii - pad_y
                        jjj = jj - pad_x
                        new_value += kernel[ii, jj] * data_part[i + iii, j + jjj, k]
                data_part_new[i - pad_y, j - pad_x, k] = new_value
    return data_part_new


def split_array(arr, n, overlap):
    sub_array_length = (len(arr) + (n - 1) * overlap) // n
    sub_arrays = [arr[i * (sub_array_length - overlap): i * (sub_array_length - overlap) + sub_array_length] for i in range(n)]
    return np.array(sub_arrays)


def process_image(img, func=0, mode=0):
    if isinstance(img, str):
        img = Image.open(img)

    if img.mode == "P":
        img = img.convert(mode="RGB")

    data = np.asarray(img, dtype=np.float64) / 255.0

    if func == 1:
        kernel = uniform_kernel(7)
    else:
        kernel = gauss_kernel(3)

    padding = [(kernel.shape[0] // 2), kernel.shape[1] // 2]

    if mode == 1:  # Multithreading mode
        num_threads = 16

        data_parts = split_array(data, num_threads, padding[0])

        with multiprocessing.Pool(num_threads) as pool:
            data_new_parts = pool.starmap(process_image_part, zip(data_parts, [kernel]*num_threads, [padding]*num_threads))
        data_new = np.concatenate(data_new_parts, axis=0)

    elif mode == 2:  # MPI parallel mode
        comm = MPI.COMM_WORLD
        rank = comm.Get_rank()
        size = comm.Get_size()
        data_part = np.array_split(data, size, axis=0)[rank]
        data_new_part = process_image_part(data_part, kernel, padding)
        data_new_parts = comm.gather(data_new_part, root=0)
        if rank == 0:
            data_new = np.concatenate(data_new_parts, axis=0)
        else:
            data_new = None
        data_new = comm.bcast(data_new, root=0)

    else:  # Default mode
        data_new = process_image_part(data, kernel, padding)

    data_new = data_new * 255.0
    data_new = np.uint8(data_new)

    return Image.fromarray(data_new, mode="RGB")



if __name__ == '__main__':
    print(monte_carlo_methode(1000, FLAG_default))
    print(monte_carlo_methode(1000, FLAG_threaded))
    print(monte_carlo_methode(1000, FLAG_network))

    url = "https://i.wfcdn.de/teaser/660/27020.jpg"
    response = requests.get(url)

    if response.status_code == 200:
        image = Image.open(BytesIO(response.content))
        image = process_image(image, FLAG_uniform, FLAG_threaded)
        image.show()