############################################
# Researched and Implemented by: Ethan Huang
# Mentored by: Rudy Osuna
# Triton Quantitative Trading @ UC San Diego
############################################
# region imports
from AlgorithmImports import *
# endregion


class EMAOnlyCrossoverBenchmark(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5 * 365))
        self.set_cash(100_000)
        self.settings.seed_initial_prices = True
        self.universe_settings.resolution = Resolution.DAILY
        self.universe_settings.schedule.on(self.date_rules.week_start())
        self._universe = self.add_universe(self._select)
        self.set_warm_up(250)

    def on_warmup_finished(self):
        time_rule = self.time_rules.at(8, 0)
        # Rebalance the portfolio weekly after warmup.
        self.schedule.on(self.date_rules.week_start("SPY"), time_rule, self._rebalance)
        # Rebalance the portfolio today too.
        if self.live_mode:
            self._rebalance()
        else:
            self.schedule.on(self.date_rules.today, time_rule, self._rebalance)

    def _select(self, fundamental):
        # Select the top 100 most liquid stocks by dollar volume with price above $10.
        return [f.symbol for f in sorted(
            [f for f in fundamental if f.has_fundamental_data and f.price > 10],
            key=lambda f: f.dollar_volume
        )[-100:]]

    def on_securities_changed(self, changes):
        # Once a security leaves the universe, liquidate and remove indicators.
        for security in changes.removed_securities:
            if security.invested:
                self.liquidate(security.symbol)
            self.deregister_indicator(security.fast_ema)
            self.deregister_indicator(security.slow_ema)
        # Initialize EMA indicators for newly added securities.
        for security in changes.added_securities:
            security.fast_ema = self.ema(security.symbol, 50, Resolution.DAILY)
            security.slow_ema = self.ema(security.symbol, 200, Resolution.DAILY)
            self.warm_up_indicator(security.symbol, security.fast_ema, Resolution.DAILY)
            self.warm_up_indicator(security.symbol, security.slow_ema, Resolution.DAILY)

    def _rebalance(self):
        if self.is_warming_up:
            return
        selected_symbols = []
        for symbol in self._universe.selected:
            security = self.securities[symbol]
            if not (security.fast_ema.is_ready and security.slow_ema.is_ready):
                continue
            if security.fast_ema.current.value > security.slow_ema.current.value:
                selected_symbols.append(security.symbol)
        if len(selected_symbols) == 0:
            for symbol in self._universe.selected:
                security = self.securities[symbol]
                if security.invested:
                    self.liquidate(security.symbol)
            return
        weight = 1.0 / len(selected_symbols)
        targets = [PortfolioTarget(symbol, weight) for symbol in selected_symbols]
        self.set_holdings(targets, True)