from AlgorithmImports import *
from QuantConnect.DataSource import *

class ETFConstituentUniverseAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2010, 1, 1)
        self.set_end_date(2020, 1, 1)
        self.set_cash(1_000_000_000)
        self.universe_settings.asynchronous = True
        self.universe_settings.resolution = Resolution.MINUTE

        # Requesting data
        self.spy = self.add_equity("SPY").symbol 
        self._universe = self.add_universe(self.universe.etf(self.spy, self.universe_settings, self.etf_constituents_filter))
        
        # Example of Historical Universe Data
        history = self.history(self._universe, 30, Resolution.DAILY)
        for (universe_symbol, time), constituents in history.items():
            for constituent in constituents:
                self.debug(f'{constituent.symbol}: {constituent.weight}')

        self.weight_by_symbol = {}
        
        # Set up daily rebalance scheduled event
        self.schedule.on(
            self.date_rules.every_day(self.spy),
            self.time_rules.after_market_open(self.spy, 1),
            self.rebalance)

    def etf_constituents_filter(self, constituents: List[ETFConstituentUniverse]) -> List[Symbol]:
        # Get the 10 securities with the largest weight in the index, save the weights
        selected = sorted([c for c in constituents if c.weight],
            key=lambda c: c.weight, reverse=True)[:10]
        self.weight_by_symbol = {c.symbol: c.weight for c in selected}
        
        return list(self.weight_by_symbol.keys())

    def rebalance(self) -> None:
        spy_weight = sum(self.weight_by_symbol.values())

        # Liquidate the ones not in top 10 weights
        if spy_weight > 0:
            for symbol in self.portfolio.Keys:
                if symbol not in self.weight_by_symbol:
                    self.liquidate(symbol)

            # Invest portfolio by normalized weights of the top 10 constituents
            for symbol, weight in self.weight_by_symbol.items():
                self.set_holdings(symbol, 1 * weight / spy_weight) 

    def on_securities_changed(self, changes: SecurityChanges) -> None:
        for security in changes.removed_securities:
            if security.invested:
                self.liquidate(security.symbol, 'Removed From Universe')

        for security in changes.added_securities:
            # Historical data for newly added constituents
            history = self.history(security.symbol, 7, Resolution.DAILY)
            self.debug(f'We got {len(history)} prices for asset from our history request for {security.symbol}')