import pandas as pd

# Import our custom functions
from StationarityAndZScores import *

# Import the Liquid ETF Universe helper methods
from QuantConnect.Data.UniverseSelection import *

class NadionParticleAntennaArray(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 11, 1)  # Set Start Date
        self.SetCash(1000000)  # Set Strategy Cash

        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        
        self.SetBenchmark('SPY')
        
        self.SetExecution(ImmediateExecutionModel())

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

        self.UniverseSettings.Resolution = Resolution.Minute
        self.SetUniverseSelection(LiquidETFUniverse())
        
        self.AddEquity('XLE')
        self.Schedule.On(self.DateRules.EveryDay('XLE'), self.TimeRules.AfterMarketOpen('XLE', 5), self.TransformTestTrade)
    
    def TransformTestTrade(self):
        qb = self
        
        symbols = [x.Symbol for x in qb.ActiveSecurities.Values]
        # Copy and paste from research notebook
        # -----------------------------------------------------------------------------
        # Fetch history and returns
        history = qb.History(symbols, 500, Resolution.Hour)
        returns = history.unstack(level = 1).close.transpose().pct_change().dropna()
        
        # Test for stationarity
        stationarity = TestStationartiy(returns)
        
        # Get z-scores
        z_scores = GetZScores(returns)
        # -----------------------------------------------------------------------------

        insights = []

        # Iterate over symbols
        for symbol, value in stationarity.iteritems():
            # Only emit Insights for those whose returns exhibit stationary behavior
            if value:
                # Get most recent z_score
                z_score = z_scores[symbol].tail(1).values[0]
                if z_score < -1:
                    insights.append(Insight.Price(symbol, timedelta(1), InsightDirection.Up))
                elif z_score > 1:
                    if self.Portfolio[symbol].Invested:
                        insights.append(Insight.Price(symbol, timedelta(1), InsightDirection.Flat))

        self.EmitInsights(insights)

import numpy as np
import pandas as pd
from statsmodels.tsa.stattools import adfuller

def TestStationartiy(returns):
    # Return pandas Series with True/False for each symbol
    return pd.Series([adfuller(values)[1] < 0.05 for columns, values in returns.iteritems()], index = returns.columns)
    
def GetZScores(returns):
    # Return pandas DataFrame containing z-scores
    return returns.subtract(returns.mean()).div(returns.std())