import numpy as np
import pandas as pd

from statsmodels import regression
import statsmodels.api as sm
import matplotlib.pyplot as plt
import math

import pandas_datareader.data as web
from datetime import datetime

def Starter_Lec7():
    start = datetime(2014, 1, 1)
    end = datetime(2015, 1, 1)

    asset = web.DataReader("TSLA","yahoo",start,end)
    asset_closingPrice = asset['Close']

    benchmark = web.DataReader("SPY","yahoo",start,end)
    benchmark_closingPrice = benchmark['Close']
    
    r_a = asset_closingPrice.pct_change()[1:]
    r_b = benchmark_closingPrice.pct_change()[1:]
    
    modelSummary = linreg(r_a,r_b)
    print("{0} {1} \n\n".format(modelSummary,type(modelSummary)))


def linreg(X,Y):
    #running linear regression
    X = sm.add_constant(X)
    model = regression.linear_model.OLS(Y,X).fit()    

    a = model.params[0]
    b = model.params[1]

    X = pd.DataFrame(X, columns=['Close'])     #Y_CMT Neat trick to extract columns from a pandas dataframe

    # Return summary of the regression and plot results
    X2 = np.linspace(float(X.min()), float(X.max()), 100)
    Y_hat = X2 * b + a

    plt.scatter(X, Y, alpha=0.3) # Plot the raw data

    plt.plot(X2, Y_hat, 'r', alpha=0.9)  # Add the regression line, colored in red
    plt.xlabel('X Value')
    plt.ylabel('Y Value')    
    plt.show()
    return model.summary()

def TestPlotting():
    N = 8
    y = np.zeros(N)

    x1 = np.linspace(0, 10, N, endpoint=True)
    x2 = np.linspace(0, 10, N, endpoint=False)

    plt.plot(x1, y, 'o')
    plt.plot(x2, y + 0.5, 'o')
    plt.ylim([-0.5, 1])
    plt.show()


def NumpyMatrix():
    array1 = np.matrix([[1,2,3],[4,5,6],[7,8,9]])
    print("{0} {1} \n\n".format(array1[:,2],type(array1)))

    array1 = array1[:,2]
    print("{0} {1} \n\n".format(array1,type(array1)))
