# 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:]
                ]
            )
        )
        #self._universe = self.add_universe(
        #    lambda fundamentals: [f.symbol for f in sorted(fundamentals, key=lambda f: f.market_cap)[-20:]]
        #)

        # Create a Scheduled Event to record new daily prices and
        # rebelance the portfolio.
        self.schedule.on(
            self.date_rules.every_day(spy),
            # The fund is dynamically rebalanced on a quarterly basis - the 3rd Friday
            # of March, June, September, and December.
            # Source: https://www.ishares.com/us/strategies/top-20#top-5-companies
            #FuncDateRule(
            #    name="third_friday_of_the_quarter_end_month", 
            #    get_dates_function=self._third_friday_of_the_quarter_end_month
            #),
            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])
        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)

    def _third_friday_of_the_quarter_end_month(self, start_date, end_date):
        dt = start_date.replace(day=1).replace(tzinfo=None)
        dates = []
        while dt <= end_date:
            if dt.month in [3, 6, 9, 12]:
                days_until_friday = (4 - dt.weekday()) % 7  # 4 represents Friday
                first_friday = dt + timedelta(days=days_until_friday)
                third_friday = first_friday + timedelta(weeks=2)
                dates.append(third_friday)
            dt += relativedelta(months=1)
        return dates