# region imports
from AlgorithmImports import *
# endregion

class DeterminedFluorescentPinkOwlet(QCAlgorithm):

    def Initialize(self) -> None:
        
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2023, 5, 1)
        self.SetWarmup(50)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.qqq = self.AddEquity("QQQ").Symbol #ETF1
        self.swan = Symbol.Create("SPY", SecurityType.Equity, Market.USA)

        self.AddUniverse(self.Universe.ETF(self.qqq, self.UniverseSettings, self.ETFConstituentsFilter_QQQ))
        self.AddUniverse(self.Universe.ETF(self.swan, self.UniverseSettings, self.ETFConstituentsFilter_SWAN))
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.Midnight, self.Rebalance)

        self.AddRiskManagement(NullRiskManagementModel())
        self.SetExecution(ImmediateExecutionModel())
        self.Settings.FreePortfolioValuePercentage = 0.20 #cash buffer 20%
        

        self.weightBySymbol_QQQ = {}
        self.weightBySymbol_SWAN = {}

    def ETFConstituentsFilter_QQQ(self, constituents):
        selected = sorted([c for c in constituents if c.Weight],
            key=lambda c: c.Weight, reverse=True)[:10] 
        self.weightBySymbol_QQQ = {c.Symbol: c.Weight for c in selected}
        
        return list(self.weightBySymbol_QQQ.keys())
    
    def ETFConstituentsFilter_SWAN(self, constituents):
        self.Log("Triggered for SWAN")
        selected = sorted([c for c in constituents if c.Weight],
            key=lambda c: c.Weight, reverse=True)[:10] 
        self.Log(selected)
        self.weightBySymbol_SWAN = {c.Symbol: c.Weight for c in selected}
        
        return list(self.weightBySymbol_SWAN.keys())

    def Rebalance(self):
        qqqWeight = sum(self.weightBySymbol_QQQ.values())
        swanWeight = sum(self.weightBySymbol_SWAN.values())
                
        if swanWeight > 0:
            for symbol in self.Portfolio.Keys:
                if symbol not in self.weightBySymbol_SWAN:
                    self.Liquidate(symbol)
    
            for symbol, weight in self.weightBySymbol_SWAN.items():
                self.SetHoldings(symbol, 0.03 * weight / swanWeight)

        if qqqWeight > 0:
            for symbol in self.Portfolio.Keys:
                if symbol not in self.weightBySymbol_QQQ:
                    self.Liquidate(symbol)
    
            for symbol, weight in self.weightBySymbol_QQQ.items():
                self.SetHoldings(symbol, (1-0.03) * weight / qqqWeight)

    def OnSecuritiesChanged(self, changes):
        for security in changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol, 'Removed From Universe')