from sklearn import linear_model
reg = linear_model.LinearRegression()
reg.fit(ages_train,net_worths_train)


#####################
slope and intercept
#####################
print "katie's net worth prediction: ", reg.predict([27])
print "slope: ", reg.coef_
print "intercept: ", reg.intercept_


####################
####### R ##########
print "\n ####### stats on test dataset #######\n"
print "r-squared score: ", reg.score(ages_test, net_worths_test)
print "\n ####### stats on training dataset #####\n"
print "r-squared score: ", reg.score(ages_train, net_worths_train)

###################
#### visualize ####
from sklearn.linear_model import LinearRegression

reg = LinearRegression()
reg.fit(ages,net_worths)
print "Katie's net worth prediction: ", reg.prediction([27])

print "r-squared score: ", reg.score(ages, net_worths)
print "slope: ", reg.coef_
print "intercept: ", reg.intercept_


plt.scatter(ages, net_worths)
plt.plot(ages, reg.prediction(ages), color='blue',linewidth= 3)
plt.xlabel("ages")
plt.ylabel("net worths")
plt.show()