import numpy as np
import pandas as pd
import statsmodels.api as sm
from statsmodels import regression
import matplotlib.pyplot as plt
import scipy.stats
from quantopian.pipeline import Pipeline
from quantopian.pipeline.data import Fundamentals
from quantopian.pipeline.data.builtin import USEquityPricing
from quantopian.pipeline.factors import CustomFactor, DailyReturns
from quantopian.pipeline.filters import QTradableStocksUS, Q500US
from quantopian.research import run_pipeline

class Factor(CustomFactor):
    inputs = [USEquityPricing.close]
    window_length = 126

    def compute(self, today, assets, out, prices):
        out[:] = (prices[-10] - prices[-126]) / prices[-126]


start = "2019-01-01"
end = "2019-12-03"


pipe = Pipeline(
    columns={
        'Returns': DailyReturns(),
        'Factor': Factor()
    },
    screen=QTradableStocksUS()
)

df = run_pipeline(pipe, start, end)

df.head()

df.dropna(inplace=True)

B = pd.DataFrame(index=df.index.get_level_values(level=0).unique(), dtype=np.float32)

for i in B.T:
    y = df.loc[i]['Returns']
    x = df.loc[i]['Factor']
    result = sm.OLS(y, x).fit()
    
    B.loc[i,'beta'] = result.params[0]

B

B.describe()