from AlgorithmImports import *
import talib
from datetime import timedelta
from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Algorithm.Framework.Alphas import *

class TTMAlphaModel(AlphaModel):
    def __init__(self, period=20, BB_mult=2, ATR_mult = 1.5, rsi_period=14, rsi_overbought=70, rsi_oversold=30):
        self.period = period
        self.k = BB_mult
        self.ATR_mult = ATR_mult
        self.rsi_period = rsi_period
        self.rsi_overbought = rsi_overbought
        self.rsi_oversold = rsi_oversold
        self.last_close = {}
        self.bb = BollingerBands(period, BB_mult)
        self.kch = KeltnerChannels(period, ATR_mult, MovingAverageType.Exponential)
        self.atr = AverageTrueRange(period, MovingAverageType.Exponential)
        self.prev_squeeze = {}


    def Update(self, algorithm, data):
        insights = []
        
        # Get current time
        current_time = algorithm.Time


        # ZO --- WEAK slipapge modelling, but best we have available at the moment.
        # algorithm.Debug(f'Time: {current_time}')
        
        algo = algorithm 
        qb = algo.History(algo.Securities.Keys, 100, Resolution.Minute)
        for symbol, security in algorithm.Securities.items():
            # security.SetSlippageModel(VolumeShareSlippageModel())
            security.SetSlippageModel(ConstantSlippageModel(.0005)) # .05% slippage (default)
            try:
                df = qb.loc[symbol]
                spread_pct = (df.askclose.mean() - df.bidclose.mean()) / df.askclose.mean()
                security.SetSlippageModel(ConstantSlippageModel(spread_pct))
            except:
                pass
            


        # Get current universe
        universe = algorithm.UniverseManager.ActiveSecurities
        
        # Get historical data for universe
        # Calculate TTM Squeeze and RSI indicators for each security in universe
        for security in universe:
            history = algorithm.History(security.Value.Symbol, 30, Resolution.Daily)

            bar = data.Bars.get(security.Value.Symbol)

            if bar:

                self.bb.Update(bar.EndTime, bar.Close)
                self.kch.Update(bar)
                self.atr.Update(bar)

                if not history.empty and self.bb.IsReady and self.kch.IsReady and self.atr.IsReady:                  
                    # Get last close price
                    current_close = data[security.Value.Symbol].Close
                    
                    # Calculate Bollinger Bands, Keltner Channels, and True Range
                    bb_upper, _, bb_lower = talib.BBANDS(history['close'], timeperiod=self.period)
                    kama = talib.KAMA(history['close'], timeperiod=self.period)
                    # Calculate ATR
                    atr = talib.ATR(history['high'], history['low'], history['close'], timeperiod=20)
                    
                    mom = talib.MOM(history['close'], timeperiod=20)

                    if len(mom) < 5:
                        continue
                    smoothed_mom = mom.rolling(5).mean()

                    kc_upper = kama + (self.ATR_mult * atr)
                    kc_lower = kama - (self.ATR_mult * atr)

                    # Calculate TTM Squeeze
                    if bb_upper[-1] < kc_upper[-1] and bb_lower[-1] > kc_lower[-1]:
                        squeeze = True
                    else:
                        squeeze = False

                    if bb_upper[-2] < kc_upper[-2] and bb_lower[-2] > kc_lower[-2]:
                        prev_squeeze = True
                    else:
                        prev_squeeze = False                    

                    mom_bullish = smoothed_mom[-1] > smoothed_mom[-2] and smoothed_mom[-1] > 0 and smoothed_mom[-2] > 0 #Blue
                    mom_bearish = smoothed_mom[-1] < smoothed_mom[-2] and smoothed_mom[-1] < 0 and smoothed_mom[-2] < 0 #Red

                    mom_bullish_stop = smoothed_mom[-1] < smoothed_mom[-2] #Dark Blue
                    mom_bearish_Stop = smoothed_mom[-1] > smoothed_mom[-2] #Yellow                
                    
                    # # Calculate RSI
                    # rsi = talib.RSI(history['close'], timeperiod=self.rsi_period)

                    # overbought = rsi[-1] > self.rsi_overbought
                    # oversold = rsi[-1] < self.rsi_oversold
                    # stop = rsi[-1] < self.rsi_overbought and rsi[-1] > self.rsi_oversold

                    # check for TTM Squeeze and mom is momentum indicator
                    if mom_bullish:
                        if squeeze and prev_squeeze:
                            insights.append(Insight.Price(security.Value.Symbol, timedelta(30), InsightDirection.Up))
                    elif mom_bearish:                        
                        if squeeze and prev_squeeze:
                            insights.append(Insight.Price(security.Value.Symbol, timedelta(30), InsightDirection.Down))               

                    if algorithm.Portfolio[security.Value.Symbol].Invested:
                        if algorithm.Portfolio[security.Value.Symbol].IsLong and mom_bullish_stop:
                            insights.append(Insight.Price(security.Value.Symbol, timedelta(1), InsightDirection.Flat))                    
                            #algorithm.Liquidate(security.Value.Symbol.Value, "Liquidated exit short")

                        elif algorithm.Portfolio[security.Value.Symbol].IsShort and mom_bearish_Stop:
                            insights.append(Insight.Price(security.Value.Symbol, timedelta(1), InsightDirection.Flat))
                            #algorithm.Liquidate(security.Value.Symbol.Value, "Liquidated exit short")                                            

                    # Update last_close
                    self.last_close[security] = current_close
                    #self.prev_squeeze[security.Value.Symbol] = squeeze
        
        return insights
        
class TTMAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 3, 15)
        self.SetEndDate(2023, 4, 1)
        self.SetCash(100000)
        self.maxDD_security = 0.15
        
        # Universe selection
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction)
        
        # Alpha model
        self.SetAlpha(TTMAlphaModel())
        
        # Portfolio construction and risk management
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(self.maxDD_security))
        self.Settings.RebalancePortfolioOnInsightChanges = True
        self.Settings.RebalancePortfolioOnSecurityChanges = False

        self.Settings.Resolution = Resolution.Minute  # ZO -- not working.
        #self.UniverseSettings.ExtendedMarketHours = False

        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)


        # ZO
        # # CHANGE to minute -- force it to use minute on Data 0
        # self.symbol = self.AddEquity("SPY", Resolution.Minute).Symbol
        # # Set benchmark
        # self.SetBenchmark(self.symbol)
        
    def CoarseSelectionFunction(self, coarse):
        """
        Perform coarse filters on universe.
        Called once per day.
        Returns all stocks meeting the desired criteria.
        
        Attributes available:
         .AdjustedPrice
         .DollarVolume
         .HasFundamentalData
         .Price -> always the raw price!
         .Volume
        """
        # Get the highest volume stocks
        stocks = [x for x in coarse if x.HasFundamentalData]
        sorted_by_dollar_volume = sorted(
            stocks, key=lambda x: x.DollarVolume, reverse=True
        ) 
        top = 50
        symbols = [x.Symbol for x in sorted_by_dollar_volume[:top]]
        # Print universe details when live mode
        if self.LiveMode:
            self.MyLog(f"Coarse filter returned {len(symbols)} stocks.")
        return symbols