#!/usr/bin/env python3

import numpy as np
from scipy.integrate import solve_ivp
from PIL import Image
import colorsys
import math
from concurrent.futures import ProcessPoolExecutor
import json

# this function was made by GPT because I am way too fucking stupid to know how to code double pendulum physics
def double_pendulum_simulation(theta1_start, theta2_start, sim_time) -> tuple[float]:
    """
    Simulates a double pendulum and returns the final angles of both pendulums.

    Parameters:
        theta1_start (float): Initial angle of the first pendulum in radians. Angle 0 is downwards and + goes counter clockwise.
        theta2_start (float): Initial angle of the second pendulum (in radians). Angle 0 is downwards and + goes counter clockwise.
        sim_time (float): Time to simulate the system (in seconds).

    Returns:
        tuple: Final angles (theta1, theta2) in radians. Angle is in the range [0 - 2pi)
    """
    # Constants (fixed values)
    g = 9.81  # gravitational acceleration (m/s^2)

    # Equations of motion
    def equations(t, y):
        theta1, omega1, theta2, omega2 = y

        delta_theta = theta1 - theta2
        denom1 = 2 - np.cos(delta_theta)**2
        denom2 = denom1

        alpha1 = (-g * (2) * np.sin(theta1)
                  - g * np.sin(theta1 - 2*theta2)
                  - 2 * np.sin(delta_theta) * (omega2**2 + omega1**2 * np.cos(delta_theta))) / denom1

        alpha2 = (2 * np.sin(delta_theta) * (omega1**2
                                             + g * np.cos(theta1)
                                             + omega2**2 * np.cos(delta_theta))) / denom2

        return [omega1, alpha1, omega2, alpha2]

    # Initial conditions
    theta1_0 = theta1_start  # initial angle of first pendulum (radians)
    omega1_0 = 0.0           # initial angular velocity of first pendulum
    theta2_0 = theta2_start  # initial angle of second pendulum (radians)
    omega2_0 = 0.0           # initial angular velocity of second pendulum

    y0 = [theta1_0, omega1_0, theta2_0, omega2_0]

    # Time span
    t_span = (0, sim_time)

    # Solve using Runge-Kutta method
    solution = solve_ivp(equations, t_span, y0, method='RK45')

    # Final angles
    theta1_final = solution.y[0, -1]
    theta2_final = solution.y[2, -1]

    return float(theta1_final % (2 * np.pi)), float(theta2_final % (2 * np.pi))


RESOLUTION: int = 4000

def process_column(x):
    column_pixels = []
    for y in range(RESOLUTION):
        angle1, angle2 = double_pendulum_simulation(2 * np.pi / RESOLUTION * x, 2 * np.pi / RESOLUTION * y, 2)

        angle1Norm: float = float(angle1 / (2 * np.pi))
        angle2Norm: float = float(angle2 / (2 * np.pi))

        # normalize angle one
        if angle1Norm <= 0.5:
            angle1Norm *= 2
        else:
            angle1Norm = 1 - ((angle1Norm - 0.5) * 2)

        # make image brighter
        angle1Norm **= (1 / 3)

        color = colorsys.hsv_to_rgb(angle2Norm, 1, angle1Norm)

        # Store pixel data for this column
        column_pixels.append((x, y, round(color[0] * 255), round(color[1] * 255), round(color[2] * 255)))

    with open(f"row_data/row_{x}", "w") as f:
        json.dump(column_pixels, f)
    print(x)


choice = input("press enter to start generating image. Write anything and then press enter to combine row data into one image file.")
if choice == "":
    # Use ProcessPoolExecutor to parallelize the task
    with ProcessPoolExecutor(max_workers=8) as executor:
        executor.map(process_column, range(RESOLUTION))
else:
    image = Image.new("RGB", (RESOLUTION, RESOLUTION))
    pixels = image.load()

    for row in range(RESOLUTION):
        with open(f"row_data/row_{row}") as f:
            columnPixels = json.load(f)

        for x, y, r, g, b in columnPixels:
            pixels[x, y] = (r, g, b)

    image.save("output.png")