# region imports
from AlgorithmImports import *
# endregion

class MeasuredTanTermite(QCAlgorithm):

    stopMarketTicket = None
    stopMarketFillTime = datetime.min

    def Initialize(self):
        #self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
        self.SetStartDate(2022, 8, 1)
        self.SetEndDate(2022, 8, 20)
        self.SetCash(1000000)
        self.rebalanceTime = datetime.min
        self.AddUniverse(self.CoarseSelectionFilter, self.FineSelectionFilter)
        self.UniverseSettings.Resolution = Resolution.Minute #resolution of the securities in the universe, default is minute
        self.AddEquity("SPY")
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 10), self.ExitPositions)
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.At(9, 31), self.ScanTargets)
        self.Data = {}
        self.AddRiskManagement(TrailingStopRiskManagementModel(0.05))
        self.open_prices = {}
        self.Targets = {}
        self.filled_today = {}
        self.myOrder = None

# Create a new array of filled daily positions, in OnOrderEvent if the filled status is true then add symbol to this array. In OnData check if this symbol is NOT in the array before being allowed to place a new order. Lastly, clear the array at EOD in the liquidate scheduled event. 
    def OnData(self, data):
        for symbol in self.Targets.keys():
            symbolData = self.Data[symbol]

            if not symbolData.IsReady: 
               continue

    def ExitPositions(self):
        for symbol in self.Targets.keys():
            self.Debug("Target Symbol before clearing: " + str(symbol))

        self.Targets = {}

    def ScanTargets(self):
        self.Targets = {}
        for symbol in self.Data.keys():
            symbolData = self.Data[symbol]
            self.Debug("Symbols in Data dict: " + str(symbol) + " " + str(len(symbol.Value)))

            if not symbolData.IsReady: 
                continue

            self.Targets[symbol] = SymbolData(self, symbol)

        for symbol in self.Targets.keys():
            self.Debug("Targets: " + str(symbol))

    def CoarseSelectionFilter(self, coarse): #updates at midnight during backtesting, and around 6-7am in live trading

        if self.Time <= self.rebalanceTime: #checks if it has been 1 day since rebalancing
            return self.Universe.Unchanged #if not do not change the universe
        self.rebalanceTime = self.Time + timedelta(1) #set rebalance time to 1 day

        sortedByDollarVolume = sorted(coarse, key=lambda c: c.DollarVolume, reverse=True)
        symbols_by_price = [c.Symbol for c in sortedByDollarVolume if c.Price > 1 and c.Price < 5 and c.HasFundamentalData and c.Volume > 50000 and len(c.Symbol.Value) <= 4]
        self.filteredByPrice = symbols_by_price[:10]
        return self.filteredByPrice

    def FineSelectionFilter(self, fine):
        sortedByMarketCap = sorted(fine, key=lambda c: c.MarketCap)
        symbols_by_marketcap = [c.Symbol for c in sortedByMarketCap if c.MarketCap > 0]
        self.filteredBymarketcap = symbols_by_marketcap[:8]
        return self.filteredBymarketcap

    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.Data:
                self.Data[symbol] = SymbolData(self, symbol)
                self.Debug("Added to universe: " + str(len(symbol.Value)) + "   " + str(symbol.Value))
        
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.Data:
                symbolData = self.Data.pop(symbol, None)
                self.SubscriptionManager.RemoveConsolidator(symbol, symbolData.consolidator)
                self.Debug("Removed from universe: " + str(len(symbol.Value)) + "   " + str(symbol.Value))

class SymbolData:
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        self.Bars = RollingWindow[TradeBar](5) # Rolling window for data bars
        self.consolidator = TradeBarConsolidator(timedelta(days=1))
        self.consolidator.DataConsolidated += self.OnDataConsolidated
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
    
    def OnDataConsolidated(self, sender, bar):
        self.Bars.Add(bar)
    
    @property
    def IsReady(self):
        return self.Bars.IsReady