from AlgorithmImports import *
from QuantConnect.DataSource import *

class ETFConstituentsDataAlgorithm(QCAlgorithm):

    def Initialize(self) -> None:
        self.SetStartDate(2016, 1, 1)
        self.SetEndDate(2021, 1, 1)
        self.SetCash(100000)
        
        self.UniverseSettings.Resolution = Resolution.Minute
        
        # Requesting data
        self.spy = self.AddEquity("SPY").Symbol
        self.AddUniverse(self.Universe.ETF(self.spy, self.UniverseSettings, self.ETFConstituentsFilter))
        
        self.weightBySymbol = {}
        
        self.Schedule.On(
            self.DateRules.EveryDay(self.spy),
            self.TimeRules.AfterMarketOpen(self.spy, 1),
            self.Rebalance)

    def ETFConstituentsFilter(self, constituents: List[ETFConstituentData]) -> List[Symbol]:
        # Get the 10 securities with the largest weight in the index
        selected = sorted([c for c in constituents if c.Weight],
            key=lambda c: c.Weight, reverse=True)[:10]
        self.weightBySymbol = {c.Symbol: c.Weight for c in selected}
        
        return list(self.weightBySymbol.keys())

    def Rebalance(self) -> None:
        spyWeight = sum(self.weightBySymbol.values())

        if spyWeight > 0:
            for symbol in self.Portfolio.Keys:
                if symbol not in self.weightBySymbol:
                    self.Liquidate(symbol)
    
            for symbol, weight in self.weightBySymbol.items():
                self.SetHoldings(symbol, 0.5 * weight / spyWeight)
                
            self.SetHoldings(self.spy, -0.5)

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

        for security in changes.AddedSecurities:
            # Historical data
            history = self.History(security.Symbol, 7, Resolution.Daily)
            self.Debug(f'We got {len(history)} from our history request for {security.Symbol}')