'''
   Using numpy's polynomial.fit() to perform a linear 
   regression on Hubble's velocity versus distance data 
   with the constant/intercept set to zero
'''

import matplotlib.pyplot as plt  
# may need to run at the terminal the command: pip install matplotlib
import numpy as np  
# may need to run at the terminal the command: pip install numpy

# declare two lists for the distances and velocities
dist =[0.032, 0.034, 0.214, 0.263, 0.275, 0.275, 0.45, 0.5, 0.5, 0.63, 
       0.8, 0.9, 0.9, 0.9, 0.9, 1, 1.1, 1.1, 1.4, 1.7, 2, 2, 2, 2]
vel = [170, 290, -130, -70,-185,-220, 200, 290, 270, 200, 300, -30, 650, 
       150, 500, 920, 450, 500, 500, 960, 500, 850, 800, 1090]

axes=plt.axes()
axes.set_facecolor("#ffffcc")    #color plot area light blue
plt.grid(alpha=0.5)              # create semi-transparent grid

plt.scatter(dist, vel, s=30, alpha=0.7, color="#000099", edgecolors="#000033")


#slope, intercept = np.polyfit(dist, vel, deg=1)
#the_fit = np.polynomial.Polynomial.fit(dist, vel, deg=1)
# deg=[1] will skip constant i.e set the intercept to zero
the_fit = np.polynomial.Polynomial.fit(dist, vel, deg=[1], domain=[], rcond=None, full=False, w=None, window=None, symbol='x')
#intercept = the_fit(0)
slope = the_fit(1)-the_fit(0)
print(the_fit)             # weirdly "scaled"(?) result -- fixed by domain=[]
print(the_fit.convert())   # more what one expects

# Use numpy to create sequence of 10 numbers ranging 
# from minimum speed to maximum speed  (i.e. make the x's )
xseq = np.linspace(min(dist), max(dist), num=10)
# Plot regression line 
plt.plot(xseq,  slope * xseq , color="#0000ff", alpha=0.5, lw=3)

# prepare linear fit equation to be displayed
fit_eqn = "y = {:4.2f} x".format(slope)
plt.text(0.25, 800, fit_eqn, fontsize=12)

# title the plot and label x and y axes
plt.title("Hubble's Data")
plt.xlabel("Distance (mega-parsec)")
plt.ylabel("Velocity (km/s)")

import os   # run operating system commands to set current working directory
os.chdir(os.path.dirname(os.path.realpath(__file__)))
# save the figure to an external .png file
plt.savefig('Hubble.pdf')
# show the plot on the screen (in a pop-up)
plt.show()