#!/usr/bin/env python

from numpy import exp, pi
from PIL import Image, ImageDraw
from random import random
import numpy as np, os, 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 PIL library is separated as much as is practical to make porting
       to different graphics libraries easier.

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

def cli(argv):
    fname = 'univ.gif' # File name, if f > 1, to store animated gif.
    f =   1 # Number of frames.
    n = 250 # Make 250x250 calculations.
    N = 400 # Image will be 400x400 pixels.
    r =  50 # ms per frame.

    prog = argv[0]
    argv = argv[1:]
    i = 0
    while i < len(argv):
        arg = argv[i]
        if arg == '-f':   # Number of frames to create.
            i += 1
            f = int(argv[i])
        elif arg == '-g': # File name of animated gif.
            i += 1
            fname = argv[i]
        elif 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 f, n, N, fname, r

def spirals(n=100, t=0):
    '''Calculate an n x n matrix of spirals.  Each complex number
    will later be scaled, colored and plotted.

    Inputs:
      n (int): unitless, resulting matrix will be n x n in size.
      t (float): unitless, used for animation.  New starting point.
    Output:
      z (complex[][]): complex numbers to be plotted.  For all elements
        e in z, |e| <= 2.
    '''

    z = np.zeros((n,n), dtype=complex)
    for i in range(n): # Number of spirals.
        ang = complex(i, t + (2*pi/n)*i) # Explore variations of this line! :-)
        zPrev = 0 # Usually will be z[i][j-1].
        for j in range(n): # Points in this spiral.
            theta = ang + zPrev
            zPrev = z[i][j] = exp(1j*theta.real) + exp(1j*theta.imag)
    return z

def mkFrame(z, N):
    '''Map each point in the complex matrix z to NxN pixel image.  For
    each complex v in z[], |v| <= 2.  Color and plot each point where
    its red is 2i, green is 2j and blue is 255-(i+j).

    Inputs:
      z (complex[][]): complex numbers to be plotted.  For all elements
        e in z, |e| <= 2.
      N (int): pixels, width & height of eventual image.
    '''

    # Create pixels p by scaling and translating matrix z.
    c = complex(N//2, N//2) # Center coords of image.
    r = N//4                # Radius of universe is 2*r because |z|<=2.
    p = c + r*z             # Translate & scale z to fill image.

    # Color image.
    im = Image.new('RGB', (N,N))
    draw = ImageDraw.Draw(im)
    n = z.shape[0]          # Retrieve size of matrix, n x n.
    for i in range(n):
        R = (i<<1)&0xff     # RGB red byte.
        for j in range(n):
            G,B = (j<<1)&0xff, ~(i+j)&0xff # RGB green & blue bytes.
            draw.point((p[i][j].real, p[i][j].imag), fill=(R,G,B))
    return im

def usage(prog):
    prog = os.path.basename(prog)
    print('Usage: %s -f frames -g gif -n calcs -N pixels -r rate' % prog)
    print()
    print('       -f int, default 1, number of frames in animation.')
    print('       -g str, default univ.gif, name of animated GIF file.')
    print('       -n int, default 250, make n x n calculations.')
    print('       -N int, default 400, make N x N image.')
    print('       -r int, default 50, ms to display each frame.')
    sys.exit(1)

def main(argv):
    fTot, n, N, fname, r = cli(argv)
    t = random()
    if fTot == 1: # Display image.
        z = spirals(n, t)
        im = mkFrame(z, N)
        im.show()
    else: # Create animated gif.
        frames = []
        for f in range(fTot):
            if f%10 == 0: print('Frame %d' % f)
            t += random()/60 # Random movt, smaller -> slower evolution.
            z = spirals(n, t)
            im = mkFrame(z, N)
            frames.append(im)
        frames[0].save(fname, save_all=True, append_images=frames[1:],
            optimize=False, duration=r, loop=0)

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