from AlgorithmImports import *


class FundamentalFactorAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)  
        # Add TLT and SPY.
        self.settings.automatic_indicator_warm_up = True
        self._tlt = self.add_equity('TLT')
        self._spy = self.add_equity('SPY')
        self._spy.sma = self.sma(self._spy, 120, Resolution.DAILY)
        # Add a universe of US Equities.
        self._universe_size = 100
        self._date_rule = self.date_rules.month_start(self._spy)
        self.universe_settings.schedule.on(self._date_rule)
        self._universe = self.add_universe(self._select_assets)
        # Define some trading paramters.
        self._formation_days = 200
        self._portfolio_size = 10
        # Add a warmup so the algorithm trades upon deployment.
        self.set_warm_up(timedelta(45))
 
    def _select_assets(self, fundamentals):
        # Select the most liquid stocks trading above $5.
        fundamentals = [f for f in fundamentals if f.has_fundamental_data and f.price > 5]
        fundamentals = sorted(fundamentals, key=lambda f: f.dollar_volume)[-200:]
        # Select the subset of stocks with the greatest EV/EBITDA.
        fundamentals = [
            f for f in fundamentals 
            if (f.valuation_ratios.ev_to_ebitda > 0 and
                f.earning_reports.basic_average_shares.three_months * f.price > 2e9)
        ]
        selected = sorted(fundamentals, key=lambda f: f.valuation_ratios.ev_to_ebitda)[-self._universe_size:]
        return [f.symbol for f in selected]

    def on_securities_changed(self, changes):
        for security in changes.added_securities:
            # Increase the Session, so we can track daily prices from n days ago.
            security.session.size = self._formation_days
            # Warm up the Session history.
            for bar in self.history[TradeBar](security, self._formation_days, Resolution.DAILY):
                security.session.update(bar)

    def on_warmup_finished(self):
        # Add a Scheduled event to rebalance the portfolio monthly.
        time_rule = self.time_rules.at(10, 0)
        self.schedule.on(self._date_rule, 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 _rebalance(self):
        # Wait until the warmup ends.
        if self.is_warming_up:
            return
        # If SPY is below its SMA, just hold TLT.
        if self._spy.price < self._spy.sma.current.value:
            self.set_holdings(self._tlt, 1, True)
            return
        # Otherwise, calculate the momentum of all assets in the universe.
        roc_by_security = {}
        for symbol in self._universe.selected:
            security = self.securities[symbol]
            if not security.session.is_ready:
                continue
            # Compare the closing price n days ago to the intraday price right now.
            start_value = security.session[-1].close
            if not start_value:
                continue
            roc_by_security[security] = security.price / start_value - 1
        if not roc_by_security:
            return
        # Form an equal-weighted portfolio of the assets with the greatest momentum.
        selected = sorted(roc_by_security, key=lambda security: roc_by_security[security])[-self._portfolio_size:]
        weight = 1 / len(selected)
        self.set_holdings([PortfolioTarget(equity, weight) for equity in selected], True)