#!/usr/bin/env python3

###############################################################################
# Static PNG plot of two jellyfish-like creatures
#
# This script computes one static frame and saves it as a PNG.
###############################################################################

import numpy as np
import matplotlib.pyplot as plt

# =============================================================================
# 1. Tunable parameters
# =============================================================================

n_points = 20_000

# Integer frame number: 1, 2, 3, ...
#
# It is converted inside jellyfish_coordinates() to:
#
#   t = frame_id * pi / 45
frame_id = 1

canvas_width = 400
canvas_height = 400

render_width_px = 1600
render_height_px = 1600

dpi = 100

background_colour = "#090909"
point_colour = "white"
point_alpha = 0.70

# In matplotlib, s is marker area in points squared.
point_size = 1.5

output_file = "jellyfish_frame.png"

# =============================================================================
# 2. Function returning x/y coordinates for a given frame
# =============================================================================

def jellyfish_coordinates(frame_id, n_points=20_000):
    """
    Compute the x/y coordinates of the two jellyfish-like creatures.

    Parameters
    ----------
    frame_id : int
        Integer frame number: 1, 2, 3, ...

        The frame number is converted to continuous time as:

            t = frame_id * pi / 45

        This mirrors the original JavaScript update:

            t += PI / 45

        once per draw call.

    n_points : int
        Number of points in the point cloud.

    Returns
    -------
    x : numpy.ndarray
        X coordinates of all points.

    y : numpy.ndarray
        Y coordinates of all points.

    jelly_id : numpy.ndarray
        Identifier for the two creatures:
            0 for even-indexed points
            1 for odd-indexed points.
    """

    # -------------------------------------------------------------------------
    # Point index
    # -------------------------------------------------------------------------
    #
    # i = 0, 1, 2, ..., n_points - 1
    #
    # These are not x-coordinates. They are indices used to generate the point
    # cloud.
    i = np.arange(n_points)

    # -------------------------------------------------------------------------
    # Static quantities derived from i
    # -------------------------------------------------------------------------
    #
    # k = 9 * cos(i / 61)
    k = 9 * np.cos(i / 61)

    # e = i / 652 - 13
    e = i / 652 - 13

    # d = mag(k, e)^2 / 89 + 1
    #
    # In p5.js:
    #
    #   mag(k, e) = sqrt(k^2 + e^2)
    #
    # Therefore:
    #
    #   mag(k, e)^2 = k^2 + e^2
    d = (k**2 + e**2) / 89 + 1

    # m = i % 2 * 3
    #
    # This alternates between 0 and 3:
    #
    #   i even -> m = 0
    #   i odd  -> m = 3
    #
    # These two values create the two intertwined jellyfish-like creatures.
    m = (i % 2) * 3

    # Optional identifier for the two creatures.
    jelly_id = i % 2

    # -------------------------------------------------------------------------
    # Time variable
    # -------------------------------------------------------------------------
    #
    # Original JavaScript:
    #
    #   t += PI / 45
    #
    # Here:
    #
    #   frame_id = 1 -> t = pi / 45
    #   frame_id = 2 -> t = 2*pi / 45
    #   etc.
    t = frame_id * np.pi / 45

    # -------------------------------------------------------------------------
    # Main deformation term
    # -------------------------------------------------------------------------
    #
    # q = 79 - e/2*sin(k) +
    #     k/d * (6 + 5*sin(sin(d*d + e/9 - t + m)))
    q = (
        79
        - e / 2 * np.sin(k)
        + k / d * (
            6 + 5 * np.sin(
                np.sin(d**2 + e / 9 - t + m)
            )
        )
    )

    # -------------------------------------------------------------------------
    # Angular term
    # -------------------------------------------------------------------------
    #
    # c = d/1.9 + cos(t - d*3 + m)/11 - t/16 + m
    c = (
        d / 1.9
        + np.cos(t - d * 3 + m) / 11
        - t / 16
        + m
    )

    # -------------------------------------------------------------------------
    # Final coordinates
    # -------------------------------------------------------------------------
    #
    # x = q * sin(c) + 200
    # y = (q + 40) * cos(c) + 200
    x = q * np.sin(c) + 200
    y = (q + 40) * np.cos(c) + 200

    return x, y, jelly_id

# =============================================================================
# 3. Compute one frame
# =============================================================================

x, y, jelly_id = jellyfish_coordinates(
    frame_id=frame_id,
    n_points=n_points
)

# =============================================================================
# 4. Plot one static PNG
# =============================================================================

fig_width_inches = render_width_px / dpi
fig_height_inches = render_height_px / dpi

fig, ax = plt.subplots(
    figsize=(fig_width_inches, fig_height_inches),
    dpi=dpi
)

fig.patch.set_facecolor(background_colour)
ax.set_facecolor(background_colour)

ax.set_xlim(0, canvas_width)
ax.set_ylim(0, canvas_height)
ax.set_aspect("equal")

ax.set_xticks([])
ax.set_yticks([])

for spine in ax.spines.values():
    spine.set_visible(False)

ax.margins(0)
plt.subplots_adjust(left=0, right=1, top=1, bottom=0)

ax.scatter(
    x,
    y,
    s=point_size,
    c=point_colour,
    alpha=point_alpha,
    linewidths=0
)

fig.savefig(
    output_file,
    dpi=dpi,
    facecolor=fig.get_facecolor(),
    bbox_inches=None,
    pad_inches=0
)

plt.close(fig)

print(f"Saved: {output_file}")