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
        
        # 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 MyMaximumDrawdownPercentPerSecurity(RiskManagementModel):
    '''
    Provides an implementation of IRiskManagementModel that limits the drawdown 
    per holding to the specified percentage
    REF: https://github.com/QuantConnect/Lean/blob/master/Algorithm.Framework/
    Risk/MaximumDrawdownPercentPerSecurity.py
    '''
    def __init__(self, algo, max_dd_pct=0.15):
        '''
        Initializes a new instance of the MaximumDrawdownPercentPerSecurity class
        Args:
            max_dd_pct: The maximum percentage drawdown allowed 
             for any single security holding
        '''
        self.algo = algo
        self.max_dd_pct = -abs(max_dd_pct)
        # self.constant_check_multiple = constant_check_multiple

    def ManageRisk(self, algorithm, targets):
        '''
        Manages the algorithm's risk at each time step
        Args:
            algorithm: The algorithm instance
            targets: The current portfolio targets to be assessed for risk

        DO NOT USE algorithm - QuantConnect.Algorithm.QCAlgorithm object
        INSTEAD USE self.algo - main.MeanReversionAlgorithm object
        '''
        algo = self.algo
        targets = []
        # Loop through securities
        for kvp in algo.Securities:
            security = kvp.Value
            symbol_object = security.Symbol
            symbol = str(symbol_object).split(" ")[0]
            if not security.Invested:
                continue
            pnl = security.Holdings.UnrealizedProfitPercent
            if pnl < self.max_dd_pct:
                # Cancel insights
                algorithm.Insights.Cancel([symbol_object])
                # Liquidate
                targets.append(PortfolioTarget(symbol_object, 0))
        return targets

class TTMAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2017, 1, 1)
        # self.SetStartDate(2020, 1, 1)
        # self.SetEndDate(2017, 12, 31)
        # 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.SetRiskManagement(MyMaximumDrawdownPercentPerSecurity(self))

        self.Settings.RebalancePortfolioOnInsightChanges = True
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        self.Settings.FreePortfolioValuePercentage = 0.05
        #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