#! /usr/bin/env python

import numpy as np

#  We will import the pyplot 
import matplotlib.pyplot as plt

#  We are going to download data from the internet so
import urllib

#  Begin by downloading the Union2 SN cosmology data from LBL, because
#  they are fun and easy.   Moreover, you may end of teaching 
#  astro one of these days, and so this is a good example.

# initialize arrays and read in data from the web
# list of names
SN_list = ['']
z_array = np.array([])
mod_array = np.array([])
moderr_array = np.array([])
f = urllib.urlopen('http://supernova.lbl.gov/Union/figures/SCPUnion2_mu_vs_z.txt')
for line in f:
    if line[0] == '#': continue    # Ignore anything that starts with a #
    SN, z, mod, moderr = line.split()
    SN_list.append(SN)
    z_array = np.append(z_array,np.float64(z))
    mod_array = np.append(mod_array,np.float64(mod))
    moderr_array = np.append(moderr_array,np.float64(moderr))   
f.close()

# First we close whatever windows we might have:
plt.close()

# Now let us plot some points
plt.plot(z_array, mod_array)

# Notice it is a mess, by default it connects the lines

plt.savefig('SN_cosmology_example_plot1.png', format="png")
temp=raw_input("hit enter to show next plot")
# We close the window with plt.close()
plt.close()

# We clearly need axes lables
plt.xlabel('z', fontsize=20)
# But we want a Greek Letter, so we can put some LaTeX syle code with the r command:
plt.ylabel(r'$\mu$', fontsize=20)
plt.title("Union2 SN Cosmology Data")
# Now we add a format string, as follows:
plt.plot(z_array, mod_array,'ro')
#  Now this looks better, we have red circles
plt.savefig('SN_cosmology_example_plot2.png', format="png")
temp=raw_input("hit enter to show next plot")
plt.close()

# Now let's plot with the error bars
plt.errorbar(z_array, mod_array, yerr=moderr_array, fmt='.', capsize=0,
    elinewidth=1.0, ecolor=(0.6,0.0,1.0), color='green' )

# And putting lables and colors we can do:
plt.xlabel(r'$z$', fontsize=20)
plt.ylabel(r'$\mu$', fontsize=20)
plt.title("Union2 SN Cosmology Data")

plt.savefig('SN_cosmology_example_plot3.png', format="png")
temp=raw_input("hit enter to show next plot")
plt.close()

# Now there are a lot of points, so let's figure our what our distribution is in z
plt.hist(z_array, 25)
# And slap a label on it
plt.xlabel(r'$z$', fontsize=20)
plt.savefig('SN_cosmology_example_plot4.png', format="png")
temp=raw_input("hit enter to show next plot")
plt.close()

###############################################################
# Plotting With Asymmetrical Error Bars #
###############################################################

# Now, let's find the real distance from the distance modulus.
# To do this we will define a function and a constant

def distance_Mpc(m,z):
    return 0.00001 * (10**(m/5)) / (1.0 + z)
c = 299792.458    # km/s

#  So the distance can be found with:
d_array = distance_Mpc(mod_array,z_array)

# Now we will calculate the error bars in the distance, both ways:
d_error_plus  = ( distance_Mpc((mod_array+moderr_array),z_array) - d_array )
d_error_minus = ( d_array - distance_Mpc((mod_array-moderr_array),z_array) )

# And plot the graph with asymetrical horizontal error bars, and labels. 
#  Notice the different color formats that can be used.

plt.close()
plt.errorbar(z_array, d_array, yerr=(d_error_minus,d_error_plus), fmt='s',
    capsize=5, elinewidth=1.0, color=(0.4,0.0,1.0), ecolor='aqua', barsabove=True)
plt.xlabel('z', fontsize=15, color='0.0')
plt.ylabel('Distance (Mpc)', fontsize=15, color='g')
plt.title("Union2 SN Cosmology Data", color=(0.4,0.0,1.0))

plt.savefig('SN_cosmology_example_plot5.png', format="png")
temp=raw_input("hit enter to show next plot")

###########################
# ADD A SECOND AXIS 
###########################

# First we read the existing plot limits into an array

axes1_range = np.array( plt.axis() )

# We want to use the same x-axis so we use the function.

plt.twinx()   #  This swaps the Y axis
plt.ylabel('Distance in Billions of Light Years', fontsize=15, color='purple')

# Now we will make the second axes the right scale.

axes2_range = axes1_range.copy()  # Don't forget, we need to make a copy of it.
axes2_range[2:4] = 3.26*axes1_range[2:4]/1000.0  # set second y-axis to Gly  
plt.axis(axes2_range)

plt.savefig('SN_cosmology_example_plot6.png', format="png")
temp=raw_input("hit enter to show next plot")

#################################################
# Plotting With a Multiple Axes and Annotation
#################################################

# Now imagine we want to make a plot for the public, 
# so we need to find the distance in light years. Also, we want to plot cz.

plt.close()
plt.errorbar(3.26*d_array/1000.0, c*z_array,
    xerr=(3.26*d_error_minus/1000.0, 3.26*d_error_plus/1000.0),
    fmt='.', capsize=0, elinewidth=1.0, color=(0.4,0.0,1.0), ecolor='aqua',
    barsabove=False, zorder=2)
plt.ylabel('cz (km/s)', fontsize=15, color='aqua')  # plot the labels
plt.xlabel('Distance (Billions of Light Years)', fontsize=15, color='aqua')

# Now, you want to make the plot also include z and Mpc, 
# so we need to convert our axes:

axes1_range = np.array( plt.axis() ) # get the existing axes and convert to an array
axes2_range = axes1_range.copy()  # Don't forget, we need to make a copy of it.
axes2_range[0:2] = 1000*axes1_range[0:2]/3.26  # set second x-axis to Mpc  
axes2_range[2:4] = axes1_range[2:4]/c # set second y-axis to z

# Now we switch to the second axis

plt.subplots_adjust(right=0.875, top=0.8)  # make room for each axis
plt.twinx()   #  This swaps the Y axis
plt.ylabel('z', fontsize=15, color='r')  #  I am not sure why this has to be before plt.twiny()
plt.twiny()   #  This swaps the X axis 
plt.xlabel('Distance (Mpc)', fontsize=15, color='purple')
plt.title("Union2 SN Cosmology Data", color=(0.4,0.0,1.0), x=0.5, y=1.15)
plt.axis(axes2_range)

# And we will put Hubble's law on it just for fun. 
# Using the zorder keyword, we can make sure the line is plotted over 
# the points.

H = 70.8  # km/s/Mpc
z_limits = np.array([0.0,1.0])  #  Plot is relative to the new axes
d_limits = z_limits*c/H
plt.plot(d_limits,z_limits, 'y-', zorder=4)

# Notice that the second plot rescaled the axes, 
# which is a problem for this plot, 
# because both x axes and both y axes represent the same physical quantity. 
# So we reissue the axis function:

plt.axis(axes2_range)  # this needs to be after the plot command

#  Now we will put an arrow on the graph and a label.

plt.annotate("Hubble's law", xy=(d_limits[1],z_limits[1]), color='y', zorder=1, 
    xytext=(1.2*d_limits[1],1.2*z_limits[1]), arrowprops=dict(facecolor='y', shrink=0.05))
    
plt.savefig('SN_cosmology_example_plot7.png', format="png")
temp=raw_input("hit enter to show next plot")