# region imports
from AlgorithmImports import *

from dateutil.relativedelta import relativedelta
# endregion


class TOPTAnalysisAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2014, 12, 31)
        self.set_cash(1_000_000)

        spy = Symbol.create('SPY', SecurityType.EQUITY, Market.USA)
        self.set_security_initializer(
            BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices))
        )

        # Add a universe of daily data.
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(
            self.universe.etf(
                Symbol.create('SPY', SecurityType.EQUITY, Market.USA), 
                universe_filter_func=lambda constituents: [
                    c.symbol for c in sorted(
                        [c for c in constituents if c.weight], 
                        key=lambda c: c.weight
                    )[-20:]
                ]
            )
        )

        # Create a Scheduled Event to record new daily prices and
        # rebelance the portfolio.
        self.schedule.on(
            self.date_rules.every_day(spy),
            self.time_rules.at(0, 1),
            self._rebalance
        )

        self.set_warm_up(timedelta(30))

    def _rebalance(self):
        if self.is_warming_up or not self._universe.selected:
            return
        symbols = [s for s in self._universe.selected if s in self.securities and self.securities[s].price]
        if not symbols:
            return
        market_cap_sum = sum([self.securities[s].fundamentals.market_cap for s in symbols])
        self.plot('Market Cap', 'Sum', market_cap_sum)
        if not market_cap_sum:
            return
        self.set_holdings([PortfolioTarget(symbol, self.securities[symbol].fundamentals.market_cap/market_cap_sum) for symbol in symbols], True)