# region imports
from AlgorithmImports import *
from datetime import timedelta
#from QuantConnect.Data.Custom.CBOE import *
# endregion

class CombinedAlgorithm(QCAlgorithm):

    def Initialize(self):
        # Initialize
        self.SetStartDate(2023, 1, 10)
        self.SetEndDate(2023, 2, 3)
        self.SetCash(10000) 
        self.SPY = self.AddEquity('SPY', Resolution.Minute).Symbol
        self.SetWarmUp(timedelta(days=20))

        # Universe Selection
        self.AddUniverse(self.MyCoarseFilterFunction)
        self.UniverseSettings.Resolution = Resolution.Minute

        # Scheduled Events
        self.Schedule.On(self.DateRules.EveryDay(),self.TimeRules.AfterMarketOpen(self.SPY),self.MarketOpen)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.BeforeMarketClose(self.SPY, 3), self.LiquidateToggle)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.BeforeMarketClose(self.SPY, 1), self.ResetWatchlist)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.BeforeMarketClose(self.SPY), self.MarketClose)

        # Variables 
        self.Watchlist = []
        self.macdBySymbol = {}
        self.highs = {}
        self.lows = {}
        self.entry_price = {}
        self.trade_count = {}

        # Toggles 
        self.liquidate = False
        
    
    def MyCoarseFilterFunction(self, coarse):
        sorted_by_dollar_volume = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 10 and x.Price < 500], 
                                key=lambda x: x.DollarVolume, reverse=True)
        selected = [x.Symbol for x in sorted_by_dollar_volume[:20]]

        return selected
    
    # SCHEDULED FUNCTIONS
    def MarketOpen(self):
        for symbol in self.Watchlist:
            self.highs[symbol] = []
            self.lows[symbol] = []
            self.entry_price[symbol] = []
            self.trade_count[symbol] = 0
            self.Log(symbol)
    
    def LiquidateToggle(self):
        self.liquidate = True

    def ResetWatchlist(self):
        self.Watchlist = []
        self.macdBySymbol = {}
        self.liquidate = False
    
    def MarketClose(self): 
        return

    def OnData(self, data): 
        if self.IsWarmingUp:
            return   

        # Accessing Data
        release = data.Get(EstimizeRelease)

        for key, value in release.items(): 
            self.Watchlist.append(value.Symbol)

        for symbol in self.Watchlist:
            if not symbol.Value in [x.Value for x in self.macdBySymbol]: continue
            symbol_m =  [x for x in self.macdBySymbol if symbol.Value== x.Value][0]
            symbol_data = self.macdBySymbol[symbol_m]

            if symbol_data is None or not self.macdBySymbol[symbol_m].warmed_up:
                continue

            # VARIABLES
            symbol_price = self.Securities[symbol].Price
            held_stocks = self.Portfolio[symbol].Quantity
            min_high = self.Securities[symbol].High
            min_low = self.Securities[symbol].Low

            prev_macd = symbol_data.prev_macd
            current_macd = symbol_data.current_macd
            prev_macd_slope = symbol_data.prev_macd_slope
            current_macd_slope = symbol_data.current_macd_slope
            current_signal = symbol_data.current_signal

            # Twenty - Five Minute Range
            if symbol in self.highs and symbol in self.lows and len(self.highs[symbol]) < 25 and len(self.lows[symbol]) < 25:
                self.highs[symbol].append(min_high)
                self.lows[symbol].append(min_low)
            if symbol in self.highs and symbol in self.lows and len(self.highs[symbol]) == 25 and len(self.lows[symbol]) == 25:
                five_min_high = max(self.highs[symbol])
                five_min_low = min(self.lows[symbol])
                five_min_range = five_min_high - five_min_low

                
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            if security.Symbol not in self.macdBySymbol:
                estimize_release_symbol = self.AddData(EstimizeRelease, security.Symbol).Symbol
                history = self.History([estimize_release_symbol], 10, Resolution.Daily)

                self.AddEquity(security.Symbol, Resolution.Minute)
                symbol_data = SymbolData(self, security.Symbol, 12, 26, 9, MovingAverageType.Exponential, Resolution.Minute)
                self.macdBySymbol[security.Symbol] = symbol_data

                symbol_data.warmed_up = True

class SymbolData:
    def __init__(self, algorithm, symbol, fastPeriod, slowPeriod, signalPeriod, movingAverageType, resolution):
        self.symbol = symbol
        self.macd = MovingAverageConvergenceDivergence(fastPeriod, slowPeriod, signalPeriod, movingAverageType)
        self.warmed_up = False
        
        self.prev_macd = 0
        self.current_macd = 0

        self.prev_macd_slope = 0
        self.current_macd_slope = 0

        self.current_signal = 0

        # Create a 5-minute consolidator
        self.consolidator = TradeBarConsolidator(timedelta(minutes=5))
 
        # Register the consolidator with the algorithm
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)

        # Update the MACD indicator with the 5-minute bars
        self.consolidator.DataConsolidated += self.OnFiveMinuteBar

        history = algorithm.History(symbol, 1000, Resolution.Minute).loc[symbol]
        for idx, bar in history.iterrows():
            tradeBar = TradeBar(idx, symbol, bar.open, bar.high, bar.low, bar.close, bar.volume)
            self.consolidator.Update(tradeBar)
        self.warmed_up = True

    def OnFiveMinuteBar(self, sender, bar):
        self.macd.Update(IndicatorDataPoint(bar.Time, bar.Close))
        
        if self.macd.IsReady:
            self.prev_macd = self.current_macd
            self.current_macd = self.macd.Current.Value

            self.prev_macd_slope = self.current_macd_slope
            self.current_macd_slope = self.current_macd - self.prev_macd

            self.current_signal = self.macd.Signal.Current.Value