#!/usr/bin/python3
# -*- coding: utf8 -*-

# Python script from http://commons.wikimedia.org/wiki/User:Geek3/Cosmos-animation
# Creates animations of galaxies in a universe evolving according to the http://en.wikipedia.org/wiki/Friedmann_equations

try:
    import cairo as C
except ImportError:
    print('You need to install the cairo library')
    exit(1)

try:
    import scipy.optimize as op
except ImportError:
    print('You need to install SciPy: http://scipy.org/Download')
    exit(1)

import os
import random
from scipy.integrate import odeint
from math import *


width, height = 220, 160
cx, cy = (width - 1) / 2, (height - 1) / 2
bgcolor = (1, 1, 1)
circlecolor = (0.3, 0.3, 0.3)
radius = 5


models = []

def a_static_universe(t):
    # Steady state solution http://en.wikipedia.org/wiki/Static_universe
    return 1.
models.append({'name':'Cosmos-animation_Static', 'a':a_static_universe,
    'steps':1, 'tmax':0, 'a0':1., 'ngalaxies':80})

def a_empty_universe(t):
    # Evolution of an empty universe, Omega=0
    return t
models.append({'name':'Cosmos-animation_Empty', 'a':a_empty_universe,
    'steps':80, 'tmax':1, 'a0':a_empty_universe(1./79.), 'ngalaxies':8000})

def a_De_Sitter(t):
    # Omega_m=0, Omega_r=0, k=0, Lambda=1, http://en.wikipedia.org/wiki/De_Sitter_universe
    return exp(t)
models.append({'name':'Cosmos-animation_De-Sitter', 'a':a_De_Sitter,
    'steps':80, 'tmax':3.5, 'a0':1., 'ngalaxies':5000})

def a_Einstein_de_Sitter(t):
    # Omega_m=1, Omega_r=0, k=0, Lambda=0, http://en.wikipedia.org/wiki/Einstein–de_Sitter_universe
    return t**(2./3.)
models.append({'name':'Cosmos-animation_Einstein-De-Sitter', 'a':a_Einstein_de_Sitter,
    'steps':80, 'tmax':1, 'a0':a_Einstein_de_Sitter(1./79.), 'ngalaxies':1500})

def a_Friedmann_open(t):
    # Omega_m<1, Omega_r=0, k=-1, Lambda=0
    psi = op.brentq(lambda p: sinh(p) - p - t, 0, 5)
    return cosh(psi) - 1.0
models.append({'name':'Cosmos-animation_Friedmann-open', 'a':a_Friedmann_open,
    'steps':80, 'tmax':2, 'a0':a_Friedmann_open(2./79.), 'ngalaxies':1500})

def a_Friedmann_closed(t):
    # Omega_m>1, Omega_r=0, k=1, Lambda=0
    psi = op.brentq(lambda p: p - sin(p) - t, 0, 2 * pi)
    return 1.0 - cos(psi)
models.append({'name':'Cosmos-animation_Friedmann-closed', 'a':a_Friedmann_closed,
    'steps':80, 'tmax':2*pi, 'a0':a_Friedmann_closed(2*pi/79), 'ngalaxies':500})

def a_Friedmann_radiation_dominated(t):
    # Omega_m=0, Omega_r=1, k=0, Lambda=0
    return sqrt(t)
models.append({'name':'Cosmos-animation_Friedmann-radiation', 'a':a_Friedmann_radiation_dominated,
    'steps':80, 'tmax':1, 'a0':a_Einstein_de_Sitter(1/79), 'ngalaxies':1000})

def a_lambda_CDM(t):
    # Omega_m=0.3, Omega_r=0, k=0, Lambda=0.7, http://en.wikipedia.org/wiki/Lambda-CDM_model
    return odeint(lambda y, t0: (.3/y + .7*y**2)**.5, 1e-8, [0, t])[1,0]
models.append({'name':'Cosmos-animation_Lambda-CDM', 'a':a_lambda_CDM,
    'steps':80, 'tmax':2, 'a0':a_lambda_CDM(2./79.), 'ngalaxies':2000})

def a_bouncing_universe(t):
    # Omega_m=0.5, Omega_r=0, k=1.3, Lambda=0.8
    return odeint(lambda y, t0: (.5/y + .8*y**2 - 1.3)**.5, 1+1e-7, [0, fabs(t-6)])[1,0]
models.append({'name':'Cosmos-animation_Bouncing', 'a':a_bouncing_universe,
    'steps':80, 'tmax':12., 'a0':1., 'ngalaxies':400})


for model in models:
    galaxies = [(width * random.random() - 0.5, height * random.random() - 0.5)
        for i in range(model['ngalaxies'])]
    
    for frame in range(model['steps']):
        if (model['steps'] > 1):
            t = frame / (model['steps'] - 1) * model['tmax']
        else:
            t = model['tmax']
        a = model['a'](t) / model['a0']
        
        # create an image
        surf = C.ImageSurface(C.FORMAT_RGB24, width, height)
        ctx = C.Context(surf)

        # background
        ctx.new_path()
        ctx.set_source_rgb(*bgcolor)
        ctx.rectangle(0, 0, width, height)
        ctx.fill()
        
        if (a < 1e-6):
            ctx.new_path()
            ctx.set_source_rgb(*circlecolor)
            ctx.rectangle(0, 0, width, height)
            ctx.fill()
        else:
            # draw individual galaxies
            for x0, y0 in galaxies:
                x = cx + (x0 - cx) * a
                y = cy + (y0 - cy) * a
                ctx.new_path()
                ctx.set_source_rgb(*circlecolor)
                ctx.arc(x, y, radius, 0, 2*pi)
                ctx.fill()  # stroke current path

        # save as PNG
        ofname = model['name'] + '_{:0>2d}'.format(frame)
        surf.write_to_png(ofname + '.png')
        os.system('convert ' + ofname + '.png ' + ofname + '.gif && rm ' + ofname + '.png')
    os.system('gifsicle -d5 -l0 --colors 256 --careful ' + model['name'] + '_*.gif > ' + model['name'] + '.gif')
    os.system('rm ' + model['name'] + '_*.gif')