# region imports
from AlgorithmImports import *
# endregion

# -------------------------------------------------------------------
SMA = 63;
lookback = 12;
quantile = 0.1;
# -------------------------------------------------------------------

class FocusedBlackDog(QCAlgorithm):

    def initialize(self):
        self.SetStartDate(2024, 1, 1)  # Set Start Date

        self.set_security_initializer(lambda security: security.set_margin_model(SecurityMarginModel.NULL))
        self.portfolio.set_positions(SecurityPositionGroupModel.NULL);

        #self.aapl = self.AddEquity("AAPL", Resolution.Daily)
        #self.ibm = self.AddEquity("IBM", Resolution.Daily)
        self.spy = self.AddEquity("SPY", Resolution.Daily)
        #self.nvda = self.AddEquity("NVDA", Resolution.Daily)
        #self.tsla = self.AddEquity("TSLA", Resolution.Daily)
        #self.msft = self.AddEquity("MSFT", Resolution.Daily)
        #self.amzn = self.AddEquity("AMZN", Resolution.Daily)
        #self.goog = self.AddEquity("GOOG", Resolution.Daily)
        #self.nflx = self.AddEquity("NFLX", Resolution.Daily)
        #self.fb = self.AddEquity("FB", Resolution.Daily)

        self.add_universe(self.CoarseSelectionFunction)

        self.symbols=[]
        self.adv = {}
        self.monthly_rebalance=False
        self.SymbolDataBySymbol={}

    def OnEndOfDay(self):
        security_exchange_hours = self.spy.exchange.hours
        next_day = security_exchange_hours.get_next_trading_day(self.time)

        if next_day.month != self.Time.month:
            self.monthly_rebalance = True
        return

    def CoarseSelectionFunction(self, coarse):

        if self.monthly_rebalance:

            for sec in coarse:
                if sec not in self.adv:
                    # initiate a new instance of the class
                    self.adv[sec.Symbol] = AverageDollarVolume(SMA)

                    # warm up
                    history = self.History(sec.Symbol, SMA, Resolution.Daily)
                    for bar in history.iloc[:-1].itertuples():  # leave the last row
                        self.adv[sec.Symbol].Update(bar.close, bar.volume)

                # Update with newest dollar volume
                self.adv[sec.Symbol].Update(sec.AdjustedPrice, sec.Volume)

            # Sort by Average Dollar Volume
            sortedBySMADV = sorted(self.adv.items(), key=lambda x: x[1].Value, reverse=True)

            # Return new list of symbols with ADV > $100,000
            self.symbols = [x[0] for x in sortedBySMADV if x[1].Value > 100000][:500]

            return self.symbols
        return self.symbols

    def on_securities_changed(self, changes): 

        added_symbols = [ x.Symbol for x in changes.AddedSecurities ]
        self.log("{}: {}".format(self.time, changes))

        history = self.History(
            added_symbols, 
            lookback*25, 
            Resolution.Daily)

        if history.empty: return

        for symbol in added_symbols:
        ## Create SymbolData objects for any new assets
            SymbolData = MonthlyData(symbol, lookback)
            self.SymbolDataBySymbol[symbol] = SymbolData
            SymbolData.RegisterIndicators(self)
            SymbolData.WarmUpIndicators(history.loc[symbol], self)
            self.Log('{}: {}'.format(symbol,self.SymbolDataBySymbol[symbol].Return))
        
        removed_symbols = [ x.Symbol for x in changes.RemovedSecurities ]
        for removed in removed_symbols:
        ## Unsubscribed consolidator
            SymbolData = self.SymbolDataBySymbol.pop(removed, None)
            if SymbolData is not None:
                #SymbolData.RemoveConsolidators(self)
                #self.symbolDataBySymbol=self.symbolDataBySymbol.drop(removed,axis=1)
                self.liquidate(removed)

    def OnData(self,data): 
        if not self.monthly_rebalance: 
            return

        sortedByMOM = sorted(self.SymbolDataBySymbol.items(), key=lambda x: x[1].Return, reverse=True)
        self.top_quantile = sortedByMOM[:int(len(sortedByMOM)*quantile)]
        self.bottom_quantile = sortedByMOM[-int(len(sortedByMOM)*quantile):]

        to_liquidate = set(sortedByMOM) ^ (set(self.top_quantile) | set(self.bottom_quantile))
        to_liquidate = [x[0] for x in to_liquidate]

        for i in to_liquidate:
            self.liquidate(i)
            self.log('LIQUIDATE: {}'.format(i))

        #MOM = [x[0] for x in sortedByMOM]
        #for i in MOM:
        #    self.log(" MOM: {}".format(i))
        
        #top = [x[0] for x in self.top_quantile]
        #for i in top:
        #    self.log(" TOP: {}".format(i))

        #bottom = [x[0] for x in self.bottom_quantile]
        #for i in bottom:
        #    self.log(" BOTTOM: {}".format(i))

        weight = 1/len(self.top_quantile)

        top = [x[0] for x in self.top_quantile]
        for i in top: 
            self.set_holdings(i,weight)
            self.log('BUY: {}'.format(i))
        
        bottom = [x[0] for x in self.bottom_quantile]
        for i in bottom: 
            self.set_holdings(i, -1*weight)
            self.log('SELL: {}'.format(i))

        self.monthly_rebalance=False

class AverageDollarVolume(PythonIndicator):  # Average Dollar Volume
    def __init__(self, SMA):
        self.dv = RollingWindow[float](SMA)
        self.Value = 0
    
    def Update(self, price, volume):
        self.dv.Add(price * volume)

        if self.dv.IsReady:
            self.Value = pd.Series(self.dv).mean()
            return True   
            
        return False

class MonthlyData: 
    def __init__(self, symbol, lookback):
        self.Symbol = symbol
        self.Closing = CumulativeReturn(lookback)
        self.Consolidator = TradeBarConsolidator(Calendar.Monthly)
        self.previous = None
    
    def RegisterIndicators(self, algorithm):
        algorithm.register_indicator(self.Symbol, self.Closing, self.Consolidator)

    def RemoveConsolidators(self, algorithm):
        if self.Consolidator is not None:
            algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.Consolidator)

    def WarmUpIndicators(self, history, algorithm):
            for index, tuple in history.iterrows():
                tradeBar = TradeBar()
                tradeBar.Close = tuple['close']
                tradeBar.Open = tuple['open']
                tradeBar.High = tuple['high']
                tradeBar.Low = tuple['low']
                tradeBar.Volume = tuple['volume']
                tradeBar.Time = index
                tradeBar.Symbol = self.Symbol
                self.Consolidator.Update(tradeBar)
            algorithm.Log("Closing is ready after warmup: " + str(self.Closing.IsReady))
    
    @property
    def Return(self):
        return float(self.Closing.Current.Value)

    @property
    def CanEmit(self):
        if self.previous == self.Closing.Samples:
            return False

        self.previous = self.Closing.Samples
        return self.Closing.IsReady

class CumulativeReturn(PythonIndicator):  # 12-Month Cumulative Return
    def __init__(self, lookback):
        self.Indicator = RollingWindow[float](lookback)
        self.SimpleReturn = {}
        self.Value = 0

    def Update(self, bar):
        self.Indicator.Add(bar.close)

        if self.Indicator.IsReady:
            self.SimpleReturn = pd.Series(self.Indicator)
            self.SimpleReturn = self.SimpleReturn[::-1]
            self.SimpleReturn = self.SimpleReturn.pct_change().dropna()
            self.SimpleReturn = (self.SimpleReturn + 1).cumprod() - 1
            self.Value = self.SimpleReturn.iloc[-1]
            return True

        return False