import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

redshift = np.genfromtxt('sn_data.txt',skip_header= True, usecols=(0))
distance = np.genfromtxt('sn_data.txt',skip_header= True, usecols=(1))
sigma = np.genfromtxt('sn_data.txt',skip_header= True, usecols=(2))

plt.errorbar(redshift,distance, yerr = sigma, marker= '.', c = 'green', ls = 'None')
plt.xlabel('Redshift')
plt.ylabel('Distance')
plt.title('Distance versus Redshift');

hubble = 70
c = 2.99*10**5
import scipy.integrate as integrate

def lum_dist(z,omegalight = .7,hubble=70,c=2.9979*10**5):
    omegadark = 1-omegalight
    #function = (((omegalight*(1.0+x)**3.0+omegadark)**(1/2))**-1.0)
    result = integrate.quad(lambda x: (((omegalight*(1.0+x)**3.0+omegadark)**(1/2))**-1.0),0,z)[0]
    return (1+z)*(c/hubble)* result


lum_dist_vec = np.vectorize(lum_dist)


lum_dist_vec(redshift,.7);

model1 = lum_dist_vec(redshift,0)
model2 = lum_dist_vec(redshift,.3)
model3 = lum_dist_vec(redshift,1)

plt.errorbar(redshift,distance, yerr = sigma, marker= '.', c = 'green', ls = 'None')
plt.xlabel('Redshift')
plt.ylabel('Distance')
plt.title('Distance versus Redshift');
plt.plot(redshift, model1, c='red',label='Model 1')
plt.plot(redshift, model2, c='yellow',label='Model 2')
plt.plot(redshift, model3, c='blue',label='Model 3')
plt.legend();

def chi_square(x,x_error, y):
    return np.sum(((x-y)**2)/((x_error)**2))

chi_square(distance,sigma,model1)

chi_square(distance,sigma,model2)

chi_square(distance,sigma,model3)

test = np.linspace(0,1,21)
values = []
results = []
for thing in test:
    values.append (lum_dist_vec(redshift,thing))
for value in values:
    results.append (chi_square(distance,sigma,value))
show(results)

test[results.index(min(results))]

plt.scatter(test,results);
plt.xlabel('$Omega_m$');
plt.ylabel( '$\chi^2$');
plt.title('$Omega_m$ versus  $\chi^2$');

An omega_m of .15 fits this model the best
An omega_dark of .85 fits this model best
This model suggests a non-zero value of dark energy.