# Import libraries
import numpy as np
from statsmodels import regression
import statsmodels.api as sm
import matplotlib.pyplot as plt
import math
from quantopian.pipeline.data.builtin import USEquityPricing
from quantopian.pipeline import Pipeline
from quantopian.research import run_pipeline
import matplotlib.pyplot as plt
import pandas as pd
from pandas import DataFrame as df
from quantopian.pipeline.data.psychsignal import stocktwits
from quantopian.pipeline.factors import SimpleMovingAverage
import time
import datetime

from quantopian.pipeline.experimental import QTradableStocksUS

def linreg(X,Y):
    # Running the linear regression
    X = sm.add_constant(X)
    model = regression.linear_model.OLS(Y, X).fit()
    a = model.params[0]
    b = model.params[1]
    r = model.rsquared
    X = X[:, 1]

    # Return summary of the regression and plot results
    X2 = np.linspace(X.min(), X.max(), 100)
    Y_hat = X2 * b + a
    plt.figure();
    plt.plot(X, Y) # Plot the raw data
    plt.plot(X2, Y_hat, 'y');  # Add the regression line
    plt.plot(X2, (Y_hat - np.std(Y_hat)), 'g') # upper deviation
    plt.plot(X2, (Y_hat + np.std(Y_hat)), 'r') # lower deviation
    plt.xlabel('X Value')
    plt.ylabel('Y Value')
    return b

def make_pipeline():
    base_universe = QTradableStocksUS()
    close_price = USEquityPricing.close.latest
    
    # Calculate 30 day average of bull_minus_bear scores
    sentiment_score = SimpleMovingAverage(
        inputs=[stocktwits.bull_minus_bear],
        window_length=30,
    )
    
    sentiment = (sentiment_score.top(100) | sentiment_score.bottom(100))

    return Pipeline(columns={'longs': sentiment_score.top(100)}, screen=(base_universe & sentiment))

START_DATE = '2018-01-01'
END_DATE = '2018-01-03'

QTU_pipeline = run_pipeline(make_pipeline(), START_DATE, END_DATE, chunksize=252)

longs = QTU_pipeline[QTU_pipeline['longs']].index.drop_duplicates().tolist()
longs