import tkinter as tk

from tkinter import ttk





# Anti-Aliasing Rays:

# AA^2



# Light Rays:

# AA^2 * (L^2 * L_count)



# Diffuse Rays:

# AA^2 * (D^2 + L^2 * L_count) + AA^2 * (B_D + L_count * B_D)



# Specular Rays:

# AA^2 * S^2 + AA^2 * (B_S + L_count * B_S)



# Transmission Rays:

# AA^2 * T^2 + AA^2 * (B_T + L_count * B_T)



# Sub-Surface Scattering Rays:

# AA^2 * SSS^2  (No bounces for SSS)



# Volume Rays:

# AA^2 * V^2 + AA^2 * (B_V + L_count * B_V)



# Total Rays:

# Sum of all the above components







# Function to calculate total amount of rays per pixel

def calc_total(*args):

    try:

        aa_samples = int(entries["aa"].get())

        light_samples = int(entries["light"].get())

        num_lights = int(entries["num_lights"].get())



        # List of types of samples and their corresponding depth

        types_samples = ["diffuse", "specular", "transmission", "sss", "volume"]

        types_depth = ["diffuse_depth", "specular_depth", "transmission_depth", "volume_depth"]



        total = pow(aa_samples, 2)

        for t_sample, t_depth in zip(types_samples, types_depth):

            sample_val = int(entries[t_sample].get())

            depth_val = int(entries[t_depth].get())

            total += pow(aa_samples * sample_val, 2) * depth_val



        total += pow(aa_samples * light_samples, 2) * num_lights



        # Display total in the total_rays label

        total_rays.set(str(total))



        # Change background color of labels based on total ray count

        if total < 3000:

            total_bg_label.configure(background='green')

            total_label.configure(background='green')

        elif total < 15000:

            total_bg_label.configure(background='orange')

            total_label.configure(background='orange')

        else:

            total_bg_label.configure(background='red')

            total_label.configure(background='red')



    # If any of the fields contain non-integer values, display an error and reset background colors

    except ValueError:

        total_rays.set("Invalid Input")

        total_bg_label.configure(background='white')  # reset background to white if input is invalid

        total_label.configure(background='white')



# Function to increment or decrement field values using Up and Down arrow keys

def adjust_value(event):

    current_value = entries[event.widget.var].get()

    if current_value.isdigit():

        if event.keysym == "Up":

            entries[event.widget.var].delete(0, tk.END)

            entries[event.widget.var].insert(0, str(int(current_value) + 1))

        elif event.keysym == "Down" and int(current_value) > 0:

            entries[event.widget.var].delete(0, tk.END)

            entries[event.widget.var].insert(0, str(int(current_value) - 1))



# Create and configure root window

root = tk.Tk()

root.title("Arnold Rays Calculator")

root.resizable(False, False)  # make window not resizable



# Configure tkinter styles

style = ttk.Style(root)

style.configure("TLabel", font=("Arial", 14))

style.configure("Header.TLabel", font=("Arial", 18))

style.configure("Total.TLabel", font=("Arial", 36))



# Define variables and their display names

variables = ["aa", "diffuse", "specular", "transmission", "sss", "volume", "diffuse_depth", "specular_depth", "transmission_depth", "volume_depth", "light", "num_lights"]

pretty_names = ["AA Samples", "Diffuse Samples", "Specular Samples", "Transmission Samples", "SSS Samples", "Volume Samples", "Diffuse Depth", "Specular Depth", "Transmission Depth", "Volume Depth", "Light Samples", "Number of Lights"]



entries = {}  # Dictionary to store Entry widgets and their variable names

default_values = ["2"] + ["1"] * (len(variables) - 1)  # Default field values

total_rays = tk.StringVar()  # Variable to display total ray count

total_rays.trace("w", calc_total)  # Call calc_total function whenever total_rays changes



# Create an Entry widget for each variable

for var, pretty, default in zip(variables, pretty_names, default_values):

    if var == "diffuse_depth" or var == "light":

        ttk.Separator(root, orient=tk.HORIZONTAL).pack(side=tk.TOP, fill=tk.X, padx=10, pady=20)



    frame = ttk.Frame(root, padding="10 10 10 10")

    frame.pack(side=tk.TOP, fill=tk.X)

    label = ttk.Label(frame, text=f"{pretty}:", style="Header.TLabel", width=20)

    label.grid(column=0, row=0, sticky=tk.W)



    entry = ttk.Entry(frame, width=10, font=("Arial", 14))

    entry.var = var

    entry.bind("<Up>", adjust_value)

    entry.bind("<Down>", adjust_value)

    entry.grid(column=1, row=0, sticky=(tk.W, tk.E))

    entry.insert(0, default)

    entry.bind("<KeyRelease>", calc_total)

    entries[var] = entry



# Create a frame and label to display the total ray count

ttk.Separator(root, orient=tk.HORIZONTAL).pack(side=tk.TOP, fill=tk.X, padx=10, pady=20)

total_frame = ttk.Frame(root)

total_frame.pack(side=tk.TOP, fill=tk.X, expand=True, padx=10, pady=10)

total_bg_label = tk.Label(total_frame, background='white')

total_bg_label.pack(fill=tk.BOTH, expand=True)

total_label = tk.Label(total_bg_label, textvariable=total_rays, font=("Arial", 36), background="white")

total_label.pack(side=tk.TOP)



calc_total()  # Calculate and display total ray count initially



root.mainloop()  # Start the tkinter event loop