from AlgorithmImports import *
from QuantConnect.Indicators import Indicator, BarIndicator
from QuantConnect.Python import PythonSlice
from QuantConnect.Orders import OrderTicket, OrderStatus, OrderResponse, OrderEvent
import datetime as dt
import time as tm

class SymbolData:
    def __init__(self, algorithm: QCAlgorithm, symbol: Symbol):
        self.algorithm = algorithm
        self.symbol = symbol

        # Create an indicator
        self.indicators: Dict[str, IndicatorBase] = {}
        
        # price window rolling
        self.trade_bars: RollingWindow[TradeBar] = RollingWindow[TradeBar](10)
        
        # indicators
        self.indicators['ema9'] = SimpleMovingAverage(symbol.Value, 9)
        self.indicators['ema21'] = SimpleMovingAverage(symbol.Value, 21)
        self.indicators['ema50'] = SimpleMovingAverage(symbol.Value, 50)
        self.indicators['sar'] = ParabolicStopAndReverse(symbol.Value, 0.02, 0.2)
        self.indicators['adx'] = AverageDirectionalIndex(symbol.Value, 14)
        self.indicators['std'] = StandardDeviation(symbol.Value, 20)
        
        # trade targets
        self.orders: Dict[Symbol, OrderTicket] = {'stop': None, 'tp': None, 'entry': None}
        
        # warm up indicators
        for _, indicator in self.indicators.items():
            algorithm.WarmUpIndicator(symbol, indicator)

        # Create a consolidator to update the indicator
        self.consolidator = TradeBarConsolidator(1) 
        self.consolidator.DataConsolidated += self.OnDataConsolidated

        # Register the consolidator to update the indicator
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)    
        

    def OnDataConsolidated(self, sender: object, consolidated_bar: TradeBar) -> None:
        for _, indicator in self.indicators.items():
            if issubclass(type(indicator), Indicator):
                indicator.Update(consolidated_bar.Time, consolidated_bar.Close)
            elif issubclass(type(indicator), BarIndicator):
                indicator.Update(consolidated_bar)
            else:
                indicator.Update(consolidated_bar.Time, consolidated_bar.Close)
        self.trade_bars.Add(consolidated_bar)

    # If you have a dynamic universe, remove consolidators for the securities removed from the universe
    def dispose(self) -> None:
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)
    
class Test(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2021, 10, 7)  # Set Start Date
        self.SetEndDate(2022, 11, 27)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        self.SetBrokerageModel(BrokerageName.QuantConnectBrokerage, AccountType.Margin)    
        
        # list of symbols
        self.symbol_data_by_symbol: Dict[Symbol,SymbolData] = {}
        
        # Add SPY to the algorithm
        self.AddEquity("SPY", Resolution.Minute, extendedMarketHours=True)

    def OnSecuritiesChanged(self, changes: SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            self.symbol_data_by_symbol[security.Symbol] = SymbolData(self, security.Symbol)

        # If you have a dynamic universe, track removed securities
        for security in changes.RemovedSecurities:
            symbol_data = self.symbol_data_by_symbol.pop(security.Symbol, None)
            if symbol_data:
                symbol_data.dispose()
    
    def OnOrderEvent(self, orderEvent: OrderEvent) -> None:
        cur_sym_port = self.symbol_data_by_symbol[orderEvent.Symbol]
        cur_orders = cur_sym_port.orders
        if orderEvent.Status == OrderStatus.Filled:
            if cur_orders['stop']!= None and orderEvent.OrderId == cur_orders['stop'].OrderId:
                res = cur_orders['tp'].Cancel()
                if not res.IsSuccess:
                    self.Debug(f"Failed to cancel take profit order: {res}")
                else:
                    cur_orders['tp'] = None
                    cur_orders['stop'] = None
                    cur_orders['entry'] = None
            
            elif cur_orders['tp']!= None and orderEvent.OrderId == cur_orders['tp'].OrderId:
                if cur_orders['stop'].Quantity - orderEvent.FillQuantity == 0:
                    cur_orders['stop'].Cancel()
                    cur_orders['tp'] = None
                    cur_orders['stop'] = None
                    cur_orders['entry'] = None
                else:
                    nq = cur_orders['stop'].Quantity - orderEvent.FillQuantity
                    res = cur_orders['stop'].UpdateQuantity(nq)
                    if not res.IsSuccess:
                        self.Debug(f"Failed to update stop loss order: {res}")
                    shares = int(0.5*nq) if int(0.5*nq) > 5 else nq
                    cur_orders['tp'] = self.LimitOrder(orderEvent.Symbol, shares, (orderEvent.FillPrice + (orderEvent.FillPrice - cur_orders['entry'].AverageFillPrice)))
        self.symbol_data_by_symbol[orderEvent.Symbol].orders = cur_orders             

    def OnData(self, data: PythonSlice):
        sym: Symbol
        for sym in data.Keys:
            # current bar
            sym_data: TradeBar = data.get(sym)
            if sym_data == None:
                continue
            
            # data
            bars = self.symbol_data_by_symbol[sym].trade_bars
            inds = self.symbol_data_by_symbol[sym].indicators
            
            # validating indicator state
            ready = all([ind.IsReady for _, ind in inds.items()]) and bars.IsReady
            if not ready:
                continue

            # functional code
            mas = 1 if ((inds['ema9'].Current.Value > inds['ema21'].Current.Value) and (inds['ema21'].Current.Value > inds['ema50'].Current.Value)) else 0
            mas = -1 if ((inds['ema9'].Current.Value < inds['ema21'].Current.Value) and (inds['ema21'].Current.Value < inds['ema50'].Current.Value)) else mas
            
            adx = 1 if ((inds['adx'].Current.Value > 25) and (inds['adx'].PositiveDirectionalIndex.Current.Value > inds['adx'].NegativeDirectionalIndex.Current.Value)) else 0
            adx = -1 if ((inds['adx'].Current.Value > 25) and (inds['adx'].PositiveDirectionalIndex.Current.Value < inds['adx'].NegativeDirectionalIndex.Current.Value)) else adx
            
            sar = 1 if inds['sar'].Current.Value < bars[0].Close else 0
            sar = -1 if inds['sar'].Current.Value > bars[0].Close else sar
            
            bo = 1 if ((bars[1].Low < bars[2].Low) and (bars[2].Low < bars[3].Low) and (bars[0].Close > bars[1].High)) else 0
            bo = -1 if ((bars[1].High > bars[2].High) and (bars[2].High > bars[3].High) and (bars[0].Close < bars[1].Low)) else bo
            
            
            timeh = int(str(sym_data.EndTime)[11:13])
            timem = int(str(sym_data.EndTime)[14:16])
            time = 1 if (timeh>=9) and timeh<16 else 0
            time = 0 if (timeh==9 and timem<30) else time
            time = 2 if (timeh==16 and timem==0) else time

            # make a new long trade setting stopout and take profit
            if mas==1 and adx==1 and sar==1 and bo==1 and time==1:
                if not self.Portfolio[sym].Invested:
                    stopout = min([bars[0].Low, bars[1].Low, bars[2].Low, bars[3].Low])
                    shares = int(25/(1.1*(bars[0].Close-stopout)))
                    sym_port = self.symbol_data_by_symbol[sym]
                    sym_port.orders['entry'] = self.MarketOrder(sym, shares)         
                    
                    tm.sleep(5)
                    
                    if sym_port.orders['entry'].Status != OrderStatus.Filled:
                        res = sym_port.orders['entry'].Cancel()
                        if res.IsSuccess != OrderResponse.Success:
                            self.Debug(f"Failed to cancel order: {res}")
                    else:
                        sym_port.orders['stop'] = self.StopMarketOrder(sym, -1*sym_port.orders['entry'].QuantityFilled, stopout)
                        sym_port.orders['tp'] = self.LimitOrder(sym, int(-0.5*sym_port.orders['entry'].QuantityFilled), bars[0].Close+(sym_port.orders['entry'].AverageFillPrice-stopout))
            
            if (self.Portfolio[sym].Invested and not (mas==1 and adx==1 and sar==1)) or time==2:
                self.Liquidate(sym)
                self.symbol_data_by_symbol[sym].orders['stop'] = None
                self.symbol_data_by_symbol[sym].orders['tp'] = None           
                self.symbol_data_by_symbol[sym].orders['entry'] = None