# region imports
from AlgorithmImports import *
from datetime import timedelta
# endregion

class PriceMomentumStrategy(QCAlgorithm):
    """Strategy 3.1: Price-Momentum from Kakushadze's 151 Trading Strategies
    
    Each month:
    1. Select top 500 US equities by dollar volume
    2. Rank by 12-month return (excluding most recent month)
    3. Long top decile (50 stocks), short bottom decile (50 stocks)
    4. Equal-weighted, rebalance monthly
    """

    def initialize(self):
        self.set_start_date(1998, 1, 1)
        self.set_cash(1000000)
        
        self.universe_settings.resolution = Resolution.DAILY
        self.settings.seed_initial_prices = True
        
        self._universe = self.add_universe(self._select_coarse)
        
        self._num_coarse = 500
        self._num_long = 50
        self._num_short = 50
        
        self._rebalance_time = self.date_rules.month_start()
        self.schedule.on(self._rebalance_time, self.time_rules.after_market_open("SPY", 30), self._rebalance)
        
        self._momentum_period = 252
        self._skip_period = 21
    
    def _select_coarse(self, coarse):
        """Select top 500 US equities by dollar volume"""
        filtered = [x for x in coarse if x.has_fundamental_data and x.price > 5]
        
        sorted_by_volume = sorted(filtered, key=lambda x: x.dollar_volume, reverse=True)
        
        return [x.symbol for x in sorted_by_volume[:self._num_coarse]]
    
    def _calculate_momentum(self, symbol):
        """Calculate 12-month momentum excluding most recent month
        
        Returns price change from 12 months ago to 1 month ago
        """
        history = self.history(symbol, self._momentum_period + self._skip_period, Resolution.DAILY)
        
        if history.empty or len(history) < self._momentum_period + self._skip_period:
            return None
        
        prices = history.loc[symbol]['close']
        
        price_recent = prices.iloc[-self._skip_period]
        price_old = prices.iloc[0]
        
        if price_old == 0:
            return None
        
        return (price_recent - price_old) / price_old
    
    def _rebalance(self):
        """Monthly rebalancing: rank by momentum and create long/short portfolio"""
        active_securities = list(self._universe.selected)
        
        if len(active_securities) == 0:
            return
        
        momentum_data = []
        for symbol in active_securities:
            momentum = self._calculate_momentum(symbol)
            if momentum is not None:
                momentum_data.append((symbol, momentum))
        
        if len(momentum_data) < self._num_long + self._num_short:
            return
        
        sorted_by_momentum = sorted(momentum_data, key=lambda x: x[1], reverse=True)
        
        long_symbols = [x[0] for x in sorted_by_momentum[:self._num_long]]
        short_symbols = [x[0] for x in sorted_by_momentum[-self._num_short:]]
        
        targets = []
        
        long_weight = 1.0 / self._num_long
        for symbol in long_symbols:
            targets.append(PortfolioTarget(symbol, long_weight))
        
        short_weight = -1.0 / self._num_short
        for symbol in short_symbols:
            targets.append(PortfolioTarget(symbol, short_weight))
        
        self.set_holdings(targets, liquidate_existing_holdings=True)