import tkinter as tk
from tkinter import Scale
from PIL import Image, ImageTk
from matplotlib.figure import Figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import numpy as np


def calculate_knee_parameters(sensor_output1, mapped_output1, sensor_output2, mapped_output2):
    # Map inputs from [0, 255] to [0, 4095]
    in1 = (sensor_output1 << 4) + (sensor_output1 >> 4)
    out1 = (mapped_output1 << 4) + (mapped_output1 >> 4)
    in2 = (sensor_output2 << 4) + (sensor_output2 >> 4)
    out2 = (mapped_output2 << 4) + (mapped_output2 >> 4)

    return in1, out1, in2, out2


def apply_knee_to_array(image_array, in1, out1, in2, out2):
    # LUT mapping based on the knee function
    lut = np.zeros(4096, dtype=np.uint16)
    for i in range(4096):
        if i <= in1:
            lut[i] = int(i / in1 * out1)
        elif i <= in2:
            lut[i] = int((i - in1) / (in2 - in1) * (out2 - out1) + out1)
        else:
            lut[i] = int((i - in2) / (4095 - in2) * (4095 - out2) + out2)
    lut = np.clip(lut, 0, 4095)

    # Map pixel values using the LUT
    adjusted_array = lut[(image_array * 4095).astype(int)] / 4095.0
    return adjusted_array


def update_image():
    global processed_image_label, processed_image, knee_canvas_dark, knee_canvas_light

    # Get current slider values
    sensor_output1 = sensor_output1_slider.get()
    mapped_output1 = mapped_output1_slider.get()
    sensor_output2 = sensor_output2_slider.get()
    mapped_output2 = mapped_output2_slider.get()

    # Calculate knee parameters
    in1, out1, in2, out2 = calculate_knee_parameters(sensor_output1, mapped_output1, sensor_output2, mapped_output2)

    # Apply knee function
    adjusted_array = apply_knee_to_array(image_array, in1, out1, in2, out2)
    processed_image = Image.fromarray((adjusted_array * 255).astype(np.uint8))
    processed_image_resized = processed_image.resize((processed_image.width // 4, processed_image.height // 4))

    # Update the processed image in the GUI
    processed_image_tk = ImageTk.PhotoImage(processed_image_resized)
    processed_image_label.configure(image=processed_image_tk)
    processed_image_label.image = processed_image_tk

    # Calculate and update average gray levels
    source_avg_gray = np.mean(image_array) * 255
    processed_avg_gray = np.mean(adjusted_array) * 255
    source_avg_gray_label.config(text=f"Average Gray Level: {source_avg_gray:.2f}")
    processed_avg_gray_label.config(text=f"Average Gray Level: {processed_avg_gray:.2f}")

    # Update knee function graphs
    x = np.linspace(0, 4095, 4096) / 16  # Scale x-axis to 0-255
    y = np.zeros_like(x)

    for i in range(len(x)):
        if x[i] <= in1 / 16:
            y[i] = x[i] / (in1 / 16) * (out1 / 16)
        elif x[i] <= in2 / 16:
            y[i] = (x[i] - (in1 / 16)) / ((in2 - in1) / 16) * ((out2 - out1) / 16) + (out1 / 16)
        else:
            y[i] = (x[i] - (in2 / 16)) / ((4095 - in2) / 16) * ((255 - (out2 / 16))) + (out2 / 16)

    # Plot dark knee graph
    knee_ax_dark.clear()
    knee_ax_dark.plot(x, y, label="Knee Function", color="blue")
    knee_ax_dark.axvline(in1 / 16, color="red", linestyle="--", label="Knee Point 1")
    knee_ax_dark.axvline(in2 / 16, color="green", linestyle="--", label="Knee Point 2")
    knee_ax_dark.set_title("Knee Function")
    knee_ax_dark.set_xlabel("Input (0-255)")
    knee_ax_dark.set_ylabel("Output (0-255)")
    knee_ax_dark.set_xlim(0, 255)
    knee_ax_dark.set_ylim(0, 255)
    knee_ax_dark.legend()
    knee_canvas_dark.draw()


if __name__ == "__main__":
    root = tk.Tk()
    root.title("Knees Demo Tool V1.1")
    root.iconbitmap(default='Resources/icon1.ico')
    root.tk.call("source", "Resources/azure.tcl")
    root.tk.call("set_theme", "dark")

    # Load the image and convert it to a grayscale NumPy array
    input_image_path = "Resources/image.png"
    image = Image.open(input_image_path).convert("L")
    image_array = np.asarray(image, dtype=np.float32) / 255.0

    # Original image display
    tk.Label(root, text="Source", font=("Arial", 14)).grid(row=0, column=0, padx=10, pady=10)
    original_image_resized = image.resize((image.width // 4, image.height // 4))
    original_image_tk = ImageTk.PhotoImage(original_image_resized)
    original_image_label = tk.Label(root, image=original_image_tk)
    original_image_label.grid(row=1, column=0, padx=10, pady=10)
    source_avg_gray_label = tk.Label(root, text="Average Gray Level: 0.00", font=("Arial", 12))
    source_avg_gray_label.grid(row=2, column=0)

    # Processed image display
    tk.Label(root, text="Processed", font=("Arial", 14)).grid(row=0, column=1, padx=10, pady=10)
    processed_image = image
    processed_image_resized = processed_image.resize((processed_image.width // 4, processed_image.height // 4))
    processed_image_tk = ImageTk.PhotoImage(processed_image_resized)
    processed_image_label = tk.Label(root, image=processed_image_tk)
    processed_image_label.grid(row=1, column=1, padx=10, pady=10)
    processed_avg_gray_label = tk.Label(root, text="Average Gray Level: 0.00", font=("Arial", 12))
    processed_avg_gray_label.grid(row=2, column=1)

    # Dark and light sliders
    Brighten_Shadows = tk.Label(root, text="Brighten Shadows")
    Brighten_Shadows.grid(row=3, column=0, pady=5, sticky="ew")

    sensor_output1_slider = Scale(root, from_=0, to=255, orient="horizontal", label="Sensor Output 1", command=lambda x: update_image())
    sensor_output1_slider.set(64)
    sensor_output1_slider.grid(row=4, column=0, padx=10, pady=10)

    mapped_output1_slider = Scale(root, from_=0, to=255, orient="horizontal", label="Mapped Output 1", command=lambda x: update_image())
    mapped_output1_slider.set(128)
    mapped_output1_slider.grid(row=5, column=0, padx=10, pady=10)

    Darken_Bright = tk.Label(root, text="Darken Bright Areas")
    Darken_Bright.grid(row=3, column=1, pady=5, sticky="ew")

    sensor_output2_slider = Scale(root, from_=0, to=255, orient="horizontal", label="Sensor Output 2", command=lambda x: update_image())
    sensor_output2_slider.set(192)
    sensor_output2_slider.grid(row=4, column=1, padx=10, pady=10)

    mapped_output2_slider = Scale(root, from_=0, to=255, orient="horizontal", label="Mapped Output 2", command=lambda x: update_image())
    mapped_output2_slider.set(255)
    mapped_output2_slider.grid(row=5, column=1, padx=10, pady=10)

    # Knee function graph
    figure_dark = Figure(figsize=(6, 4), dpi=100)
    knee_ax_dark = figure_dark.add_subplot(111)
    knee_canvas_dark = FigureCanvasTkAgg(figure_dark, root)
    knee_canvas_dark.get_tk_widget().grid(row=6, column=0, columnspan=2, padx=10, pady=10)

    hendo_label = tk.Label(root, text="Made by Hendo 19/12/2024")
    hendo_label.grid(row=7, column=0, columnspan=2, pady=10, sticky="ew")


    # Run the GUI
    update_image()
    root.mainloop()
