from pandas import DataFrame, date_range
from quantopian.pipeline import CustomFactor, Pipeline
from quantopian.pipeline.data.builtin import USEquityPricing
from quantopian.pipeline.factors import AverageDollarVolume, Returns, SimpleMovingAverage
from quantopian.research import run_pipeline

# Simple slicing example -- first create a Returns factor, then extract the single column corresponding
# to AAPL. This creates a "Slice" object (called `returns_aapl` here) which is used as an input to a
# custom factor.
returns = Returns(window_length=30)
returns_aapl = returns[symbols(24)]

class UsesReturns(CustomFactor):
    window_length = 5
    inputs = [returns, returns_aapl]
    
    def compute(self, today, assets, out, returns, returns_aapl):
        # Print the shape of each input. Our AAPL slice `returns_aapl` should have a shape of (5, 1)
        # because our window length is 5, and by definition we only have 1 column.
        print 'Returns {0}:\n{1}'.format(returns.shape, returns)
        print '\n'
        print 'AAPL Returns Slice {0}:\n{1}'.format(returns_aapl.shape, returns_aapl)

pipe = Pipeline(columns={'uses_returns' : UsesReturns()})
run_pipeline(pipe, '2016-06-01', '2016-06-01');

# Slices are unaffected by masking -- when passing a mask to a custom factor, slice inputs are unaffected,
# even if the asset with which the slice is associated would be filtered out.
adv = AverageDollarVolume(window_length=30)
no_aapl = adv.percentile_between(90, 95)

pipe = Pipeline(columns={'uses_returns' : UsesReturns(mask=no_aapl)})
run_pipeline(pipe, '2016-06-01', '2016-06-01');

# Here is AAPL's 30-day SMA in June, 2014. One can imagine that attempting to compute the correlation
# between this timeseries and another would produce nonsensical results.
sma = SimpleMovingAverage(inputs=[USEquityPricing.close], window_length=30)
pipe = Pipeline(columns={'sma' : sma})
results = run_pipeline(pipe, '2014-05-20', '2014-06-25')
results.sma.unstack()[symbols(24)]

# Attempting to add a SimpleMovingAverage factor as an input will fail with a `NonWindowSafeInput` error.
class UsesInvalidInput(CustomFactor):
    window_length = 5
    inputs = [sma]
    
    def compute(self, today, assets, out, sma):
        pass

# This will fail.
UsesInvalidInput()

# Furthermore, Slices cannot be added to a pipeline. Attempts will fail with an
# `UnsupportedPipelineOutput` error.
sma_slice = sma[symbols(24)]
pipe = Pipeline(columns={'aapl_sma' : sma_slice})  # This will fail.

# This is how a Slice of Returns of an arbitrary asset might look if they could be added as a pipeline
# column. Instead of a multi-index DataFrame with dates and assets, we only need an index of dates
# with a single value per day.
DataFrame(['some value'] * 10, columns=['Returns'], index=date_range('2016-06-01','2016-06-10'))

# `Factor.pearsonr` takes a target term, in this case `returns_aapl`, and uses it to compute rolling
# pearson correlation coefficients with the columns of another factor, in this case `returns`. This
# example computes the correlation between each stock and SPY (the market) over 100-day look backs of
# Returns. It is recommended that a mask be used when using this method as computations over every asset
# is expensive.

# NOTE: This is equivalent to doing:
# returns_corr = RollingPearsonOfReturns(
#     target=symbols(8554),
#     returns_length=30,
#     correlation_length=100,
#     mask=adv.top(500),
# )
returns = Returns(window_length=30)
returns_spy = returns[symbols(8554)]  # Creates `Slice` object of SPY's returns.
returns_corr = returns.pearsonr(
    target=returns_spy,
    correlation_length=100,
    mask=adv.top(500),
)

pipe = Pipeline(columns={'returns_corr' : returns_corr})
results = run_pipeline(pipe, '2016-06-01', '2016-06-15')
results.returns_corr.unstack().dropna(axis=1)

# Similarly, we can compute linear regressions of the returns of every asset against the returns of a
# single asset. Notice that this method returns a multi-output factor, so we have to access each output
# as its own individual factor.

# NOTE: This is equivalent to doing:
# returns_corr = RollingLinearRegressionOfReturns(
#     target=symbols(8554),
#     returns_length=30,
#     regression_length=100,
#     mask=adv.top(500),
# )
returns_regr = returns.linear_regression(
    target=returns_spy,
    regression_length=100,
    mask=adv.top(500),
)

alpha = returns_regr.alpha
beta = returns_regr.beta
corr = returns_regr.r_value

pipe = Pipeline(columns={'alpha' : alpha, 'beta': beta, 'correlation': corr})
results = run_pipeline(pipe, '2016-06-01', '2016-06-15')

# Print the results of the `beta` factor, dropping any assets with NaNs (assets that were masked out on
# each date are filled with NaNs).
results.beta.unstack().dropna(axis=1)

# Here is another example, this time using two different terms. This is computing, on each day, the
# spearman rank correlation between each stock's 30 day returns over the previous 100 days with the
# previous 100 days of VIX. Note that VIX behaves like a Slice in that each lookback window of VIX is
# a single column of data.
from quantopian.pipeline.data.quandl import yahoo_index_vix as vix
returns = Returns(window_length=30)
returns_vix_corr = returns.spearmanr(
    target=vix.close,
    correlation_length=100,
    mask=adv.top(500),
)

pipe = Pipeline(columns={'returns_vix_correlation' : returns_vix_corr})
results = run_pipeline(pipe, '2016-06-01', '2016-06-15')
results.returns_vix_correlation.unstack().dropna(axis=1)

# Finally, here is an example of passing another factor, instead of a Slice, as the target to the
# `pearsonr` method. In this case, correlations are computed asset-wise. That is, if our base factor
# is `Returns` and our target term is sentiment data, then for each asset we are calculating the
# correlation between that asset's returns over the past `correlation_length` days and that asset's
# sentiment data over the past `correlation_length` days.
from quantopian.pipeline.data.sentdex import sentiment
returns = Returns(window_length=30)
returns_sent_corr = returns.pearsonr(
    target=sentiment.sentiment_signal,
    correlation_length=100,
    mask=adv.top(500),
)

pipe = Pipeline(columns={'returns_sentiment_correlation' : returns_sent_corr})
results = run_pipeline(pipe, '2016-06-01', '2016-06-15')
results.returns_sentiment_correlation.unstack().dropna(axis=1)