#!/usr/bin/env python

from numpy import exp, pi
from pyglet.gl import *
from random import random
import numpy as np, os, pyglet, sys

'''
This program is derived from Paul Dunn's Bubble Universe.

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

The primary differences here are:

    1. Mathematics is complex rather than real.
    2. Complex calculations are separated from plotting.
    3. The pyglet library is separated as much as is practical to make
       porting to different graphics libraries easier.

Mike Markowski, mike.ab3ap@gmail.com
June 2025
'''

class PixelWindow(pyglet.window.Window):
    '''Inherit GL Window for one specialized to implement Paul Dunn's
    Bubble Universe.  An endless animation is presented on-screen.
    '''

    def __init__(self, n, N):
        '''Instantiate class.

        Inputs:
          n (int): size of n x n complex array of spirals.
          N (int): size of N x N pixel image.
        '''

        super(PixelWindow, self).__init__(N, N, "Pauls Dunn's Bubble Universe")
        self.n = n
        self.N = N
        self.t = random()

        # Create pixels and sprite.
        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):
        '''Override window event handler to draw image.  Draw Bubble Universe
        pixels to double buffer. '''
        self.sprite.draw()

    def spirals(self):
        '''Calculate an n x n matrix of spirals, then scale, color and
        plot.
        '''

        n = self.n # Make n x n complex matrix of spirals.
        N = self.N # Make N x N pixel matrix of RGB values.

        # 1. Create spirals in complex z[].
        z = np.zeros((n,n), dtype=complex)
        for i in range(n): # Number of spirals.
            ang = complex(i, self.t + (2*pi/n)*i) # Explore variations! :-)
            zPrev = 0               # Usually will be z[i][j-1].
            for j in range(n):      # Points in this spiral.
                theta = ang + zPrev # Rotate exponentials.
                zPrev = z[i][j] = exp(1j*theta.real) + exp(1j*theta.imag)
        self.t += 1/30 # Step animation, smaller -> slower evolution.

        # 2. Map spirals to pixels, z[] -> p[].
        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             # Translate & scale z to fill image.

        # 3. Calculate color and draw pixel, p[] -> pix[].
        for i in range(0, len(self.pix), 3): # Clear buffer - any better way?
            self.pix[i:i+3] = (0,0,0)
        for i in range(n):
            R = (i<<1)&0xff               # RGB red byte == 2i.
            for j in range(n):
                G = (j<<1)&0xff           # RGB green byte == 2j.
                B = ~(i+j)&0xff           # RGB blue byte == 255-(i+j).
                x = int(p[i][j].real)     # Pixel x coord.
                y = int(p[i][j].imag)     # Pixel y coord.
                k = 3*int(y*N + x)        # Pixel 2D -> 1D location.
                self.pix[k:k+3] = (R,G,B) # Draw pixel.

    def update(self, dt):
        '''Flip (draw on screen) double buffer.  Update image with pix[].'''
        self.spirals() # Create one new frame, updating self.pix[].
        self.sprite.image = pyglet.image.ImageData(
            self.N, self.N, 'RGB', self.pix)

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

    Inputs:
      argv (string[]): command line arguments.
    Outputs:
      n (int): size of n x n complex array of spirals.
      N (int): size of N x N pixel image.
      r (int): Hz, rate of frame update.
    '''

    n = 250 # Make 250x250 calculations.
    N = 400 # Image will be 400x400 pixels.
    r =  30 # 30 Hz frame rate.

    prog = argv[0]
    argv = argv[1:]
    i = 0
    while i < len(argv):
        arg = argv[i]
        if arg == '-h': # Print help.
            usage(prog)
        elif arg == '-n': # n x n complex points to calculate.
            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': # ms, frame rate.
            i += 1
            r = float(argv[i])
        else:
            usage(prog)
        i += 1
    return n, N, r

def usage(prog):
    prog = os.path.basename(prog)
    print('Usage: %s [-h] [-n calcs] [-N pixels] [-r rate]' % prog)
    print()
    print('       -h, help message.')
    print('       -n int, default 250, make n x n calculations.')
    print('       -N int, default 400, make N x N image.')
    print('       -r int, default 30, frames/second.')
    sys.exit(1)

def main(argv):
    n, N, r = cli(argv) # Calculation matrix size, image size, frame rate.
    window = PixelWindow(n, N)
    pyglet.clock.schedule_interval(window.update, 1/r) # Hz, update.
    pyglet.app.run()

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