from statsmodels.tsa.stattools import adfuller

from AlgorithmImports import *


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

from AlgorithmImports import *
from StationarityAndZScores import test_stationarity, get_z_scores


class NadionParticleAntennaArray(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(1_000_000)
        # Add framework models.
        self.set_universe_selection(LiquidETFUniverse())
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel())
        self.set_execution(ImmediateExecutionModel())
        # Add a Scheduled Event to emit insights.
        self.schedule.on(
            self.date_rules.every_day('SPY'), 
            self.time_rules.after_market_open('SPY', 5), 
            self._transform_test_trade
        )
    
    def _transform_test_trade(self):
        qb = self
        symbols = [x.symbol for x in qb.active_securities.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 = test_stationarity(returns)
        # Get z-scores.
        z_scores = get_z_scores(returns)
        # -----------------------------------------------------------------------------
        insights = []
        # Iterate over symbols
        for symbol, value in stationarity.items():
            # Only emit Insights for those whose returns exhibit stationary behavior.
            if not value:
                continue
            # 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.emit_insights(insights)