import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from astropy import units as u
from astropy import constants as con 
# Only when running in n IPython notebook:
%matplotlib inline  

# Just for good looks, no need to worry about this now.
# You are welcome to do this but it is not necessary.
c1 = cm.Reds(0.6)
c2 = cm.Blues(0.6)
c3 = cm.PuOr(0.3)
plt.rc('xtick', labelsize=13)
plt.rc('ytick', labelsize=13)
plt.rc('font', size = 12)
plt.rc('lines', linewidth=3)
plt.rc('figure', figsize=(10.0, 7.0))
plt.rcParams['axes.formatter.limits'] = (-3,3) #

dist = 3.4159 * u.km
dist

time = 2.7183 * u.h
time

dist / time

(_).to(u.km / u.s)

kmph = u.Unit(u.km / u.hour)
kms = u.Unit(u.km / u.s)
mph = u.Unit(u.imperial.mile / u.hour)
(200 * kmph).to(mph)

speeds = np.arange(100, 200) * kmph
print(speeds.to(mph))

print speeds.to(u.lightyear/u.millisecond)

type(speeds)

plt.plot(speeds, np.sqrt(speeds.to(mph)))

np.sqrt(speeds)

np.log10(speeds)

logspeeds = np.log10(speeds.to(mph).value)  # No units anymore!
print(logspeeds)
plt.plot(speeds, logspeeds)
plt.show()

from astropy.cosmology import Planck13, LambdaCDM

pl13 = Planck13
redshifts = np.linspace(0.001, 100, 100)

pl13.lookback_time(7.8)

pl13.angular_diameter_distance(redshifts)

plt.plot(redshifts, pl13.angular_diameter_distance(redshifts))

pl13.comoving_distance(redshifts).to(u.lightyear)

plt.plot(redshifts, pl13.comoving_distance(redshifts).to(u.lightyear))

pl13.angular_diameter_distance(8.4)

plt.plot(redshifts, pl13.luminosity_distance(redshifts))

redshift = 0.3
pl13.comoving_distance(redshift)

def vel_now(z):
    return pl13.H0 * pl13.comoving_distance(z)

vel_now(redshift)

def vel_em(z): 
    return pl13.H(z) * pl13.scale_factor(z) * pl13.comoving_distance(z)

vel_em(redshift)

pl13.H(redshift)

pl13.scale_factor(redshift)

pl13.comoving_distance(redshift)

plt.plot((redshifts), vel_em(redshifts), label='Velovity now')
plt.plot((redshifts), vel_now(redshifts), label='Velocity at emission')
plt.axhline(con.c.to(kms).value, color='k', ls='--')
plt.legend(loc='best')
#plt.axis((-.1, 3, 0, 4e5))
#plt.plot(redshifts, vel_em(redshifts))

MatDom = LambdaCDM(0.7, 0.3, 0)
MatDom