# Import modules for pipeline (research and algo)
from quantopian.pipeline import Pipeline, CustomFactor

# Import module to run pipeline in research
from quantopian.research import run_pipeline

# Import any needed filters and factors
from quantopian.pipeline.filters import QTradableStocksUS, Q500US, Q1500US, StaticAssets
from quantopian.pipeline.factors import SimpleMovingAverage

# Import data
from quantopian.pipeline.data.builtin import USEquityPricing

# Other imports
from scipy import stats

# Import numpy and pandas because they rock
import numpy as np
import pandas as pd

class Amount_Above_SMA(CustomFactor):
    inputs = [USEquityPricing.close]
    window_length = 50
    
    # Set window_safe to True so it can be used as an input to another factor
    window_safe = True
    
    def compute(self, today, assets, out, close_prices):
        sma_close_prices = np.nanmean(close_prices, axis=0)
        out[:] = close_prices[-1] - sma_close_prices
                

class Consecutive_Days_Above_Zero(CustomFactor):
    """
    Returns the number previous days where the input is above zero
    """
    # No defaults. Must be set when instantiated.
    # Takes a single input
    
    def compute(self, today, assets, out, values):
        # The the first row (0) in values is the earliest day
        # The last row (-1) is the latest
        # We want to use the numpy 'argmax' method to count from the latest day.
        # So, flip the values so it's in reverse order
        values_latest_to_earliest = np.flipud(values)
        
        # Now, a little trick, if 'argmax' doesn't find a value it returns 0
        # However, we want it to return the max if all are above zero
        # So, append a row to the end with zeros so it will find the max length
        asset_count = values_latest_to_earliest.shape[1]
        zeros = np.zeros(asset_count)
        zero_array_row = zeros.reshape(1,-1)
        values_latest_to_earliest = np.append(values_latest_to_earliest, zero_array_row, axis=0)
        
        # Finally, use 'argmax' to find the first value less than zero
        out[:] = np.argmax(values_latest_to_earliest<=0, axis=0)

# Pipeline definition
def make_pipeline():
    universe = StaticAssets(symbols(['AAPL', 'IBM']))

    last_close = USEquityPricing.close.latest
    sma_50 = SimpleMovingAverage(inputs=[USEquityPricing.close], window_length=50)
    pct_sma = Percent_Above_SMA(mask=universe)
    days_above_sma = Consecutive_Days_Above_Zero(inputs=[pct_sma], window_length=200)

    return Pipeline(
        columns={
            'last_close': last_close,
            'sma_50': sma_50,
            'pct_sma': pct_sma,
            'days_above_sma': days_above_sma,
        },
        screen=universe,
    )

# Select a time range to look at
period_start = '2019-10-1'
period_end = '2019-11-1'

# Pipeline execution
data = run_pipeline(
    make_pipeline(),
    start_date=period_start,
    end_date=period_end
)

data