#!/usr/bin/env python 
'''
plot animation of galaxies from SDSS spectro catalog using H0=70km/s/Mpc
http://skyserver.sdss.org/dr15/en/tools/search/sql.aspx
using
http://docs.astropy.org/en/stable/coordinates/

Author: Augustin Guyonnet
aguyonnet@fas.harvard.edu
'''

import os, sys, re
import numpy as np
import matplotlib.pyplot as pl
import toolbox as tb
import astropy
import time
from astropy.coordinates import Distance
from astropy import units as u
from astropy import cosmology
from astropy.cosmology import WMAP5, WMAP7
from astropy.coordinates import SkyCoord
import matplotlib.animation as animation


H0   = 70.         # km/s/MPc
km   = 3.24078E-20 # Mpc
H0   = H0 * km
year = 31536000    # sec per year
H0y  = year * H0
age  = 1./H0


'''
reading input file
'''
def readtuple(cat):
    objs = []; 
    fp = open( cat, "r")
    lines = fp.readlines()
    for line in lines :
        if (line[0]!='#'):
            objs.append([float(x) for x in line.split()])
    fp.close()
    return np.array(objs)


if __name__ == "__main__":
    file      = sys.argv[1]
    data      = readtuple(file)
    time      = -1
    direction = np.sign(time)
    '''
    per interval million years
    '''
    interval  = 100
    cst       = H0y * interval* 1E6
    
    steps     = int(14E9/(interval*1E6))
    inBillions= interval * 1E6 / 1E9
    print steps
    
    expansion = []
    for i in range(steps):
        expansion.append(cst*i)
    x = []
    y = []
    for i in data:
        #if ((i[0]>0.) and (i[0])<15000.):
        if ((i[0]>-15000.) and (i[0]<15000.) and (i[2]<50.)):
            a = []
            b = []
            for exp in expansion:
                a.append(i[0]*(1+direction*exp))
                b.append(i[1]*(1+direction*exp))
            x.append(a)
            y.append(b)

    print 'H0 = ', H0
    print ' 1) distance parcouru en 1 an par megaparsec = ',  H0y, ' Mpc'
    print ' 2) age de l univers = ',  age/year, ' annee'
    print '  -> checking 1*2 = ', H0y* age/year

    x=np.array(x).transpose()
    y=np.array(y).transpose()

    pl.style.use('dark_background')
    fig = pl.figure() # initialise la figure
    line, = pl.plot([],[], 'r,', label='now')
    #pl.xlim((0,1500))
    pl.xlim((-1500,1500))
    pl.ylim((-1500,1500))

    pl.xlabel('Distance (Mpc)',fontsize=18)
    pl.ylabel('Distance (Mpc)',fontsize=18)
    pl.title('The cosmic expansion of galaxies',fontsize=18)
    legend = pl.legend()
    
    Writer = animation.writers['ffmpeg']
    writer = Writer(fps=10, metadata=dict(artist='Me'), bitrate=119800)
    # fps is frame per second

    
    def animate(i): 
        line.set_data(x[i], y[i])
        val = direction* i * inBillions
        line.set_label('%s billion years'%(val))
        legend = pl.legend()
        return line 

 
    ani = animation.FuncAnimation(fig, animate, frames=steps, interval=100, repeat=False)
    ani.save('cosmic_expansion.mp4', writer=writer, dpi=160, extra_args=['-vcodec', 'libx264'])
    pl.show()