# region imports
from AlgorithmImports import *
from universe import LargeLiquidUSEquities
# endregion


class ValueAndMomentumEverywhere(QCAlgorithm):
    # Value Indicator: Book-to-market ratio
    # Momentum Indicator: 12month returns, excluding the most recent month
    
    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(1000000)
        self.settings.seed_initial_prices = True
        # Define some parameters.
        self._momentum_lookback = 12*22  # 12 months
        self._momentum_delay = 22  # 1 month
        # Add a universe of US Equities.
        self._date_rule = self.date_rules.month_start('SPY')
        self.universe_settings.schedule.on(self._date_rule)
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = []
        self.add_universe_selection(LargeLiquidUSEquities())
        # Add a warm-up period so the algorithm trades right away
        # instead of waiting for the next month.
        self.set_warm_up(timedelta(45))
    
    def on_warmup_finished(self):
        # Add a Scheduled Event to rebalance the portfolio monthly.        
        time_rule = self.time_rules.at(8, 0)
        self.schedule.on(self._date_rule, time_rule, self._rebalance)  
        # Rebalance today too.
        if self.live_mode:
            self._rebalance()
        else:
            self.schedule.on(self.date_rules.today, time_rule, self._rebalance)

    def on_securities_changed(self, changes):
        # As assets enter the universe, add an indicator to track their momentum.
        for security in changes.added_securities:
            # Setup consolidator and indicators
            security.consolidator = self.resolve_consolidator(security, Resolution.DAILY)
            mom = RateOfChange(self._momentum_lookback - self._momentum_delay)
            security.momentum_indicator = IndicatorExtensions.of(Delay(self._momentum_delay), mom)
            self.register_indicator(security, mom, security.consolidator)
            # Warm up the indicator.
            for bar in self.history[TradeBar](security, self._momentum_lookback + self._momentum_delay):
                mom.update(bar.end_time, bar.close)
            self._universe.append(security)
        # As assets leave the universe, liquidate them and stop updating their indicators.
        for security in changes.removed_securities:
            self.liquidate(security.symbol)
            self.subscription_manager.remove_consolidator(security, security.consolidator)
            self._universe.remove(security)
    
    def _rebalance(self): 
        # Gather value and momentum factor values for the assets in the universe.
        df = pd.DataFrame()
        for security in self._universe: 
            if (not security.momentum_indicator.is_ready or 
                security.fundamentals.valuation_ratios.pb_ratio in [None, 0]):
                continue
            df.loc[security, 'Value'] = 1 / security.fundamentals.valuation_ratios.pb_ratio
            df.loc[security, 'Momentum'] = security.momentum_indicator.current.value
        # Ensure some constituents are ready.
        if df.empty: 
            self.debug(f'{self.time} Consitituents not warm yet')
            self.liquidate()
            return
        # Rank the securities on their factor values; Make dollar-neutral.
        weight_by_security = df.rank() - df.rank().mean()
        weight_by_security = weight_by_security.mean(axis=1)
        # Scale weights down so the portfolio stays within leverage constraints.
        weight_by_security /= abs(weight_by_security).sum()
        # Place trades to rebalance the portfolio.
        for security, weight in weight_by_security.items():
            self.set_holdings(security, float(weight))
            

# region imports
from AlgorithmImports import *
# endregion


class LargeLiquidUSEquities(FundamentalUniverseSelectionModel):

    def __init__(self, min_price=1):
        self.min_price = min_price

    def select(self, algorithm, fundamentals):
        # Select securities with fundamental data (exclude penny stocks).
        selected = [f for f in fundamentals if f.has_fundamental_data and f.price >= self.min_price]
        sorted_by_dollar_volume = sorted(selected, key=lambda x: x.dollar_volume, reverse=True)
        filtered = [c for c in sorted_by_dollar_volume[:500]]
        # Restrict to common stock; Remove ADRs, REITs, and Financials.
        selected = [
            f for f in filtered 
            if (f.security_reference.security_type == 'ST00000001' and
                not f.security_reference.is_depositary_receipt and
                not f.company_reference.is_reit and
                f.asset_classification.morningstar_sector_code != MorningstarSectorCode.FINANCIAL_SERVICES)
        ]
        # Select the largest stocks.
        top_market_cap = sorted(selected, key=lambda x: x.market_cap, reverse=True)
        return [f.symbol for f in top_market_cap[:50]]