#!/usr/bin/python
#
# This script generates a 3D galaxy from a number of parameters and stores
# it in an array. You can modify this script to store the data in a database
# or whatever your purpose is. THIS script uses the data only to generate a
# PNG with a 2D view from top of the galaxy.
#
# The algorithm used to generate the galaxy is borrowed from Ben Motz
# <motzb@hotmail.com>. The original C source code for DOS (including a 3D
# viewer) can be downloaded here:
#
# http://bits.bristol.ac.uk/motz/tep/galaxy.html
#
# Generation parameters:

# Number of stars in the core (Example: 2000)
NUMHUB   = 5000

# Number of stars in the disk (Example: 4000)
NUMDISK  = 10000

# Radius of the disk (Example: 90.0)
DISKRAD  = 20000.0

# Radius of the hub (Example: 45.0)
HUBRAD   = 10000.0

# Number of arms (Example: 3)
NUMARMS  = 3

# Tightness of winding (Example: 0.5)
ARMROTS  = 0.5

# Arm width in degrees (Not affected by number of arms or rotations)
# Exammple: 65.0
ARMWIDTH = 65.0

# Maximum depth of arms (Example: 2.0)
MAXDISKZ = 450.0

# Maximum depth of core (Example: 16.0)
MAXHUBZ  = 3500.0

# Maximum outlier distance from arms (Example: 25.0)
FUZZ     = 25.0


# X and Y size of the created PNG
PNGSIZE    = 800

# Background color of the created PNG
PNGBGCOLOR = ( 0, 0, 0 )

# Foreground color of the created PNG
PNGCOLOR   = ( 255, 255, 255 )

# PNG frame size
PNGFRAME   = 50

# ---------------------------------------------------------------------------

import Image
import ImageDraw
import random
import math
import sys

stars = []

def generateStars():
    # omega is the separation (in degrees) between each arm
    # Prevent div by zero error:
    if NUMARMS:
        omega = 360.0 / NUMARMS
    else:
        omega = 0.0
    i = 0
    while i < NUMDISK:
        # Choose a random distance from center
        dist = HUBRAD + random.random() * DISKRAD

        # This is the 'clever' bit, that puts a star at a given distance
        # into an arm: First, it wraps the star round by the number of
        # rotations specified.  By multiplying the distance by the number of
        # rotations the rotation is proportional to the distance from the
        # center, to give curvature
        theta = ( ( 360.0 * ARMROTS * ( dist / DISKRAD ) )
        
            # Then move the point further around by a random factor up to
            # ARMWIDTH
                + random.random() * ARMWIDTH
                
            # Then multiply the angle by a factor of omega, putting the
            # point into one of the arms
                #+ (omega * random.random() * NUMARMS )
                + omega * random.randrange( 0, NUMARMS )
                
            # Then add a further random factor, 'fuzzin' the edge of the arms
                + random.random() * FUZZ * 2.0 - FUZZ
                #+ random.randrange( -FUZZ, FUZZ )
            )
            
        # Convert to cartesian
        x = math.cos( theta * math.pi / 180.0 ) * dist
        y = math.sin( theta * math.pi / 180.0 ) * dist
        z = random.random() * MAXDISKZ * 2.0 - MAXDISKZ

        # Add star to the stars array            
        stars.append( ( x, y ,z ) )

        # Process next star
        i = i + 1
    
    # Now generate the Hub. This places a point on or under the curve
    # maxHubZ - s d^2 where s is a scale factor calculated so that z = 0 is
    # at maxHubR (s = maxHubZ / maxHubR^2) AND so that minimum hub Z is at
    # maximum disk Z. (Avoids edge of hub being below edge of disk)
    
    scale = MAXHUBZ / ( HUBRAD * HUBRAD )
    i = 0
    while i < NUMHUB:
        # Choose a random distance from center
        dist = random.random() * HUBRAD
      
        # Any rotation (points are not on arms)
        theta = random.random() * 360
        
        # Convert to cartesian
        x = math.cos( theta * math.pi / 180.0) * dist
        y = math.sin( theta * math.pi / 180.0) * dist
        z = ( random.random() * 2 - 1 ) * ( MAXHUBZ - scale * dist * dist )
        
        # Add star to the stars array
        stars.append( ( x, y, z ) )
    
        # Process next star
        i = i + 1

def drawToPNG( filename ):
    image = Image.new( "RGB", ( PNGSIZE, PNGSIZE ), PNGBGCOLOR )
    draw = ImageDraw.Draw( image )
    
    # Find maximal star distance
    max = 0
    for ( x, y, z ) in stars:
        if abs(x) > max: max = x
        if abs(y) > max: max = y
        if abs(z) > max: max = z
    
    # Calculate zoom factor to fit the galaxy to the PNG size
    factor = float( PNGSIZE - PNGFRAME * 2 ) / ( max * 2 )
    for ( x, y, z ) in stars:
        sx = factor * x + PNGSIZE / 2
        sy = factor * y + PNGSIZE / 2
        draw.point( ( sx, sy ), fill = PNGCOLOR )
      
    # Save the PNG
    image.save( filename )
          
# Generate the galaxy          
generateStars()

# Save the galaxy as PNG to galaxy.png
drawToPNG( "galaxy.png" )