#!/usr/bin/env python

from numpy import exp, pi
from random import random
import numpy as np, os, sys

'''
This program is derived from X user yuruyura's algorithm posted at
https://x.com/yuruyurau/status/1226846058728177665

I first came across Paul Dunn's Bubble Universe, a BASIC port of yuruyura's
code:

https://distillery.matrixnetwork.co.uk:3004/discussion/38/bubble-universe-performance-short-basic-graphical-demo

Continuing the tradition I port it to python. :-) The primary differences
in mine are:

    1. Mathematics is complex rather than real.
    2. Calculations are separated from plotting.
    3. Complex calculations utilize in-phase/quadrature signals with phase
       modulation.  Easier than it sounds.
    4. There are two plotting choices, live animation or animated gif.

Mike Markowski, mike.ab3ap@gmail.com
Aug 2026
'''

def cli(argv):
    '''Handle command line arguments.

    Inputs:
      argv (string[]): command line arguments.
    Outputs:
      m (int): number of spirals.
      n (int): number of points in each spiral.
      N (int): size of N x N pixel image.
      r (float): Hz, frame rate (live) or basis for GIF frame duration.
      k (float): feedback intensity.
      f (int): number of frames (GIF mode only).
      gifName (str or None): if set, export a GIF instead of live view.
      t (int): speed of evolution, higher is slower.
    '''

    # Image defaults.
    N = 400 # Image will be 400x400 pixels.
    m = 250 # Number of spirals.
    n = 250 # Number of points in each spiral.
    r =  30 # 30 Hz frame rate (live), or basis for GIF frame duration.

    # Evolution defaults.
    k  =  1 # Feedback generating original Universe.
    dt = 50 # Speed of evolution, higher is slower.

    # For animated gif creation.
    f =  60 # Default number of frames.
    gifName = None # Name of gif file.

    prog = argv[0]
    argv = argv[1:]
    i = 0
    while i < len(argv):
        arg = argv[i]
        if arg == '-h': # Print help.
            usage(prog)
        elif arg == '-f': # Number of frames, GIF mode only.
            i += 1
            f = int(argv[i])
        elif arg == '-g': # Animated GIF file name.
            i += 1
            gifName = argv[i]
        elif arg == '-k': # Feedback intensity.
            i += 1
            k = float(argv[i])
        elif arg == '-m': # m spirals.
            i += 1
            m = int(argv[i])
        elif arg == '-n': # n points per spiral.
            i += 1
            n = int(argv[i])
        elif arg == '-N': # N x N image size to generate.
            i += 1
            N = int(argv[i])
        elif arg == '-r': # Hz, frame rate.
            i += 1
            r = float(argv[i])
        elif arg == '-t': # Speed of evolution.
            i += 1
            dt = int(argv[i])
        else:
            usage(prog)
        i += 1
    return m, n, N, r, k, f, gifName, dt

def mkGif(m, n, N, dt, r, k, f, fname):
    '''Calculate f frames and save them as an animated GIF.

    Inputs:
      f (int): number of frames to generate.
      fname (str): output GIF file name.
      r (float): Hz, used to derive per-frame display duration for
        the GIF (duration_ms = 1000/r), so -r means the same thing
        (frame rate) in both live and GIF modes.
    '''

    from PIL import Image # Only imported when animated gif wanted.

    t = random()
    durationMs = 1000/r
    frames = []
    for i in range(f):
        if i%10 == 0:
            print('Frame %d/%d' % (i, f))
        z = spirals(m, n, t, k)
        pix = mkPixels(z, N, m, n)
        im = Image.frombytes('RGB', (N, N), pix)
        frames.append(im)
        t += 1/dt # Step animation, bigger dt -> slower evolution.

    frames[0].save(fname, save_all=True, append_images=frames[1:],
        optimize=False, duration=durationMs, loop=0)
    print('Wrote %s (%d frames)' % (fname, f))

def mkPixels(z, N, m, n):
    '''Vectorized pixel fill.  R,G,B are built once with broadcasting,
    then scattered into a flat N*N RGB buffer by index.  Colors are
    normalized by the actual m x n span so the full 0-255 range is
    used even when m or n is small (e.g., m=10, n=20), not just when
    m,n are in the hundreds.

    Inputs:
      z (complex[m][n]): spiral points, |e| <= 2, from spirals.
      N (int): pixels, width & height of output image.
    Output:
      bytes, length 3*N*N, row-major RGB pixel buffer.
    '''

    c = complex(N>>1, N>>1) # Center coords of image.
    r = N>>2                # Radius of universe is 2*r because |z|<=2.
    p = c + r*z

    x = np.clip(p.real.astype(int), 0, N-1) # Clip for rare |z|==2 -> pixel N.
    y = np.clip(p.imag.astype(int), 0, N-1)
    pixIdx = y*N + x # Flat pixel (not byte) index, m x n.

    mDiv = max(m-1, 1) # Avoid div by zero if only 1 spiral/point requested.
    nDiv = max(n-1, 1)
    i = np.arange(m)[:, None]
    j = np.arange(n)[None, :]
    R = (255*i/mDiv).astype(np.uint8)
    G = (255*j/nDiv).astype(np.uint8)
    B = (~(255*(i + j)/(mDiv + nDiv)).astype(int)) & 0xff

    rgb = np.empty((m, n, 3), dtype=np.uint8)
    rgb[..., 0] = R
    rgb[..., 1] = G
    rgb[..., 2] = B

    pixImg = np.zeros((N*N, 3), dtype=np.uint8) # Cleared each frame.
    pixImg[pixIdx.ravel()] = rgb.reshape(-1, 3)
    return pixImg.tobytes()

def runLive(m, n, N, dt, r, k):
    '''Open a window and animate the spirals forever at r Hz.'''

    import pyglet
    from pyglet.gl import GLubyte

    class PixelWindow(pyglet.window.Window):
        '''Use GL Window for present yuruyura's algorithm in an endless
        animation on screen.
        '''

        def __init__(self, m, n, N, dt, k):
            super(PixelWindow, self).__init__(N, N,
                "yuruyura's Bubble Universe")
            self.k = k
            self.m = m
            self.n = n
            self.N = N
            self.dt = dt
            self.t = random()

            rawBytes = bytearray(3*N*N)                         # 3 bytes/pixel
            self.pix = (GLubyte*len(rawBytes))(*rawBytes)       # -> GL bytes
            img = pyglet.image.ImageData(N, N, 'RGB', self.pix) # -> image
            self.sprite = pyglet.sprite.Sprite(img)             # -> sprite!

        def on_draw(self):
            self.sprite.draw()

        def frame(self):
            '''Calculate one m x n matrix of spirals, then scale,
            color and plot.'''
            z = spirals(self.m, self.n, self.t, self.k)
            self.t += 1/self.dt # Step animation, smaller -> slower evolution.
            pix = mkPixels(z, self.N, self.m, self.n)
            self.pix = (GLubyte*len(pix))(*pix)

        def update(self, dt):
            self.frame() # Create one new frame, updating self.pix[].
            self.sprite.image = pyglet.image.ImageData(
                self.N, self.N, 'RGB', self.pix)

    window = PixelWindow(m, n, N, dt, k)
    pyglet.clock.schedule_interval(window.update, 1/r) # Hz, update.
    pyglet.app.run()

def spirals(m, n, t, k):
    '''Vectorized spiral calculation.  The m spirals are computed
    together as columns of an m x 1 matrix.

    Inputs:
      m (int): number of spirals.
      n (int): number of points in each spiral.
      t (float): unitless, new animation starting point.
      k (float): feedback intensity.
    Output:
      z (complex[m][n]): for all elements e in z, |e| <= 2.
    '''

    col = np.arange(m, dtype=float)[:, None] # Spiral index, shape m x 1.
    phi = t + 2*pi*col/m
#   phi = t + 2*pi*np.ones(m)[:,None]/m # Explore phase variations!
    a = col.copy()         # Every phasor 1 start phase, m x 1.
    b = phi.copy()         # Every phasor 2 start phase, m x 1.
    z = np.zeros((m, n), dtype=complex)
    for j in range(n):     # Recurrence, can't avoid this loop.
        P1 = np.exp(1j*a)  # Phasor 1, all m at once.
        P2 = np.exp(1j*b)  # Phasor 2, all m at once.
        Z = P1 + P2        # All m spirals.
        z[:,j] = Z[:,0]    # Save results for plotting.
        a = col + k*Z.real # All spirals, phase modulate and apply feedback k.
        b = phi + k*Z.imag
    return z

def usage(prog):
    prog = os.path.basename(prog)
    print('Usage: %s [-h] [-f frames] [-g gifFile] [-k feedback]' % prog)
    print('           [-m spirals] [-n pts] [-N pixels] [-r rate] [-t speed]')
    print('       -h, help message.')
    print('       -f int, default 60, GIF mode only: number of frames.')
    print('       -g str, if given, animated GIF file name.')
    print('       -k float, default 1, feedback intensity.')
    print('       -m int, default 250, number of spirals.')
    print('       -n int, default 250, points in each spiral.')
    print('       -N int, default 400, make N x N image.')
    print('       -r float, default 30, Hz frame rate for live or GIF.')
    print('       -t int, default 50, speed of evolution, bigger is slower.')
    sys.exit(1)

def main(argv):
    m, n, N, r, k, f, gifName, t = cli(argv)
    if gifName:
        mkGif(m, n, N, t, r, k, f, gifName)
    else:
        runLive(m, n, N, t, r, k)

if __name__ == '__main__':
    main(sys.argv)