import time

import numpy as np
from PIL import Image
import concurrent.futures


# Estimates the absolute area by using the composite trapezoidal rule.
# https://numpy.org/doc/stable/reference/generated/numpy.trapz.html
def calculate_area(f, a, b, num_points=10000):
    x = np.linspace(a, b, num_points)
    y = f(x)
    area = np.trapz(np.abs(y), x)
    return area


# Estimates the integral by using the composite trapezoidal rule.
# https://numpy.org/doc/stable/reference/generated/numpy.trapz.html
def calculate_integral(f, a, b, num_points=10000):
    x = np.linspace(a, b, num_points)
    y = f(x)
    integral = np.trapz(y, x)
    return integral


def mysterious_transformation(data):
    y_pixel, x_pixel, _ = data.shape
    data_new = np.zeros([y_pixel, x_pixel, 3], dtype=np.uint8)
    for i in range(1, y_pixel - 1):
        for j in range(1, x_pixel - 1):
            for k in range(3):
                new_value = 5 * data[i, j, k] - data[i, j - 1, k] - data[i, j + 1, k] - data[i - 1, j, k] - data[
                    i + 1, j, k]
                data_new[i, j, k] = max(0, min(new_value, 255))
    return data_new


def mysterious_transformation_parallel(data, stripes=16):
    # Divide and conquer!
    y_pixel, x_pixel, _ = data.shape
    stripe_height = y_pixel // stripes
    stripe_width = x_pixel
    stripe_arrays = [data[i * stripe_height:(i + 1) * stripe_height, :] for i in range(stripes)]

    def calculate(stripe, i):
        y_stripe_pixel, x_stripe_pixel, _ = stripe.shape
        data_new = np.zeros([y_stripe_pixel, x_stripe_pixel, 3], dtype=np.uint8)
        height = i * stripe_height
        width = i * stripe_width
        for i in range(height, y_stripe_pixel - 1):
            for j in range(width, x_stripe_pixel - 1):
                for k in range(3):
                    # No data race is to be expected (old image is only read)
                    new_value = 5 * data[i, j, k] - data[i, j - 1, k] - data[i, j + 1, k] - data[i - 1, j, k] - data[i + 1, j, k]
                    data_new[i, j, k] = max(0, min(new_value, 255))
        return data_new

    i = 0
    with concurrent.futures.ThreadPoolExecutor(stripes) as executor:
        # Compute for each slide and resample afterward
        data_new_stripes = list(executor.map(calculate, stripe_arrays, [i]))
        i += 1
    return np.concatenate(data_new_stripes)


img = Image.open("lokomotive.png")
pixels = np.asarray(img, dtype=np.uint8)

# Measure execution time for mysterious_transformation
# AMD Ryzen 7 5800X: 4.976848363876343 seconds
start_time_normal = time.time()
data_new_normal = mysterious_transformation(pixels)
end_time_normal = time.time()
execution_time_normal = end_time_normal - start_time_normal
print("Execution time for normal method:", execution_time_normal, "seconds")

# Measure execution time for mysterious_transformation_parallel
# Execution time for parallel method: 0.0123138427734375 seconds
start_time_parallel = time.time()
data_new_parallel = mysterious_transformation_parallel(pixels, 256)
end_time_parallel = time.time()
execution_time_parallel = end_time_parallel - start_time_parallel
print("Execution time for parallel method:", execution_time_parallel, "seconds")

# Each time one pixel is different - don't know the reason? But this difference is negligible
num_different_pixels = np.count_nonzero(np.sum(data_new_normal != data_new_parallel, axis=-1))
print("Number of different pixels:", num_different_pixels)

img_new_normal = Image.fromarray(data_new_normal, 'RGB')
img_new_normal.save("new_normal_lokomotive.png")

img_new_parallel = Image.fromarray(data_new_normal, 'RGB')
img_new_parallel.save("new_parallel_lokomotive.png")