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, k=2, rsi_period=14, rsi_overbought=70, rsi_oversold=30):
        self.period = period
        self.k = k
        self.rsi_period = rsi_period
        self.rsi_overbought = rsi_overbought
        self.rsi_oversold = rsi_oversold
        self.last_close = {}
        self.bb = BollingerBands(period, k)
        self.kch = KeltnerChannels(period, 1.5, MovingAverageType.Exponential)
        self.atr = AverageTrueRange(period, MovingAverageType.Exponential)
        self.prev_squeeze = {}


    def Update(self, algorithm, data):
        insights = []
        
        # Get current time
        current_time = algorithm.Time
        
        # Get current universe
        universe = algorithm.UniverseManager.ActiveSecurities
        
        # Get historical data for universe
        #history = algorithm.History(universe, self.period, Resolution.Daily)
        #history = []
        # 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:
                    #kc_upper = self.kch.UpperBand.Current.Value
                    #kc_lower = self.kch.LowerBand.Current.Value
                    #bb_upper = self.bb.UpperBand.Current.Value
                    #bb_lower = self.bb.LowerBand.Current.Value
                    #atr = self.atr.TrueRange.Current.Value                    
                    # 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)
                    #kc_upper, _, kc_lower = talib.KC(history['high'], history['low'], history['close'], timeperiod=self.period, mult=self.k)
                    #tr = talib.TRANGE(history['high'], history['low'], history['close'])
                    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()


                    ema_8 = talib.EMA(history['close'], timeperiod=8)
                    ema_21 = talib.EMA(history['close'], timeperiod=21)

                    kc_upper = kama + (1.5 * atr)
                    kc_lower = kama - (1.5 * atr)

                    # Calculate TTM Squeeze
                    #if atr[-1] < kc_upper[-1] - kc_lower[-1] and bb_upper[-1] - bb_lower[-1]> kc_upper[-1] - kc_lower[-1]:
                    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                    

                    # Calculate RSI
                    rsi = talib.RSI(history['close'], timeperiod=self.rsi_period)

                    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

                    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")                          
                    '''
                    # Check for squeeze
                    if current_close > bb_upper[-1] and current_close > kc_upper[-1]:
                        if squeeze and ema_8[-1] < ema_21[-1]:
                            insights.append(Insight.Price(security.Value.Symbol, timedelta(14), InsightDirection.Down))
                    elif current_close < bb_lower[-1] and current_close < kc_lower[-1]:                        
                        if squeeze and ema_8[-1] > ema_21[-1]:
                            insights.append(Insight.Price(security.Value.Symbol, timedelta(14), InsightDirection.Up))
                    else:
                        insights.append(Insight.Price(security.Value.Symbol, timedelta(1), InsightDirection.Flat))
                    '''
                    # Check for oversold/overbought RSI
                    '''
                    if rsi[-1] > self.rsi_overbought:
                        insights.append(Insight.Price(security.Value.Symbol, timedelta(1), InsightDirection.Down))
                    elif rsi[-1] < self.rsi_oversold:
                        insights.append(Insight.Price(security.Value.Symbol, timedelta(1), InsightDirection.Up))
                    '''
                    # 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(2017, 1, 1)
        self.SetEndDate(2023, 4, 1)
        self.SetCash(100000)
        
        # 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(0.1))
        self.Settings.RebalancePortfolioOnInsightChanges = True
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        #self.UniverseSettings.ExtendedMarketHours = False
        # Set benchmark
        #self.SetBenchmark("SPY")
        
    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