from AlgorithmImports import *
from collections import deque
import numpy as np

class GAPM(PythonIndicator):

    def __init__(self, period, signal_period):
        """
        :param period: The lookback period over which to calculate price gaps.
        :param signal_period: The period over which to apply the EMA to the Gap Ratio.
        """
        self.period = period
        self.signal_period = signal_period
        self.ema            = ExponentialMovingAverage(signal_period)        
        self.prev_close     = None
        self.gaps           = deque(maxlen=period)
        self.Value          = 0
        self.WarmUpPeriod   = max(period, signal_period)
        
    @property
    def IsReady(self) -> bool:
        return self.ema.IsReady and (len(self.gaps) == self.gaps.maxlen)

    def Update(self, input_data):
        """
        :param input_data: The input price data (bar).
        :return: True if the indicator is ready, False otherwise.
        """
        if input_data is None:
            return False

        if self.prev_close is not None:
            gap = input_data.Open - self.prev_close
            self.gaps.append(gap)

            up_gaps = sum(g for g in self.gaps if g > 0)
            dn_gaps = sum(abs(g) for g in self.gaps if g < 0)

            
            up_gap_ratio = 1 if dn_gaps == 0 else 100 * up_gaps / dn_gaps
            self.ema.Update(input_data.Time, up_gap_ratio)

            self.Current.Value    = self.ema.Current.Value
            self.Value            = self.Current.Value 
            
        self.prev_close = input_data.Close

        return self.IsReady

from AlgorithmImports import *
from GAPM import *

class GapSignals(QCAlgorithm):
    """
    Gap Momentum Trading Strategy

    This is an implementation of Perry Kaufman's Gap Momentum strategy from TASC magazine (Jan'24).
    It uses the GAPM indicator to track the ratio of cumulative upward to downward price gaps.
    The strategy enters long when the EMA of the Gap Ratio is rising and a trend filter is met.
    It exits long when the EMA of the Gap Ratio is falling.

    Modifications: 
    Added simple trend filter and used EMA instead of SMA for Gap ratio.
    
    Reference:
    https://financial-hacker.com/the-gap-momentum-system/
    https://www.traders.com/Documentation/FEEDbk_docs/2024/01/TradersTips.html#item5

    u/shock_and_awful
    """
    
    ## System method, algo entry point
    def Initialize(self):
        self.InitBacktest()
        self.InitData()
        self.InitIndicators()
    
    ## Backtest Params
    def InitBacktest(self):
        self.SetStartDate(2011, 1, 1)
        self.SetEndDate(2023, 12, 30)
        self.SetCash(100000)
        
    # Subscribe to asset feed
    def InitData(self):
        self.ticker = "QQQ"
        self.symbol = self.AddEquity(self.ticker, Resolution.Daily).Symbol
        self.SetBenchmark(self.ticker)

    ## Init Indicators
    def InitIndicators(self):
        
        # Create GAPM, EMA indicators and set warmup period
        self.gapm     = GAPM(40, 20)
        self.emaFast  = ExponentialMovingAverage(50)
        self.emaSlow  = ExponentialMovingAverage(200)        
        self.SetWarmup(200, Resolution.Daily)

        # Register Indicators w/Timeframe
        self.RegisterIndicator(self.symbol, self.gapm, Resolution.Daily)
        self.RegisterIndicator(self.symbol, self.emaFast, Resolution.Daily)
        self.RegisterIndicator(self.symbol, self.emaSlow, Resolution.Daily)

        # Initialize GAPM Rolling Window (track rise/fall over n bars)
        self.windowLength   = 2
        self.gapmWindow     = RollingWindow[float](self.windowLength)
        self.gapm.Updated  += self.OnGapmUpdated

    ## Update the rolling window when a new GAPM value is calculated
    def OnGapmUpdated(self, indicator, data):
        self.gapmWindow.Add(self.gapm.Current.Value)
        return

    ## System method, called as every new bar of data arrives
    ## Logic for entry / exit signals and trades. 
    def OnData(self, data):        

        # Make sure data is available and indicators are ready
        if (self.ticker not in data ) or (data[self.ticker] is None) \
            or not (self.gapm.IsReady and self.gapmWindow.IsReady and \
                    self.emaFast.IsReady and self.emaSlow.IsReady and not self.IsWarmingUp):
            return
        
        # If we're not invested, check for the rising gaps (and go long)
        if not self.Portfolio.Invested:
            if self.UpGapsRising() and (self.emaFast.Current.Value >= self.emaSlow.Current.Value) :            
                self.SetHoldings(self.symbol, 1)

        # If we're alread invested, check for falling gaps (and liquidate)                    
        else:
            if self.UpGapsFalling():
                self.Liquidate(self.symbol)    
    
    ## Check if gaps are rising in sequence. 
    ## Note: Rolling windows store items in reverse.
    def UpGapsRising(self):
        upGapsOrderedList   = list(self.gapmWindow)[::-1]
        upGapsRising        = all(upGapsOrderedList[i] < upGapsOrderedList[i+1] for i in range(len(upGapsOrderedList)-1))
        return upGapsRising

    ## Check if gaps are falling in sequence. 
    def UpGapsFalling(self):
        upGapsReversedList  = list(self.gapmWindow)
        upGapsFalling       = all(upGapsReversedList[i] < upGapsReversedList[i+1] for i in range(len(upGapsReversedList)-1))
        return upGapsFalling