from AlgorithmImports import *

class AdaptiveMomentumVolatility(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2023, 1, 1)
        self.set_cash(100000)
        
        # universe selection
        self.add_universe(self.coarse_selection_function, self.fine_selection_function)
        
        self.data = {}
        self.num_coarse = 50   # take 50 liquid stocks
        self.num_fine = 20     # filter down to 20
        self.max_positions = 10
        
        # ATR stop-loss
        self.atr_period = 14
        self.stop_loss_atr = 3  

        # VIX proxy (VIXY ETF since true VIX not tradable)
        self.vix = self.add_equity("VIXY", Resolution.DAILY).symbol

        # rebalance monthly
        self.schedule.on(
            self.date_rules.month_start("SPY"),
            self.time_rules.after_market_open("SPY"),
            self.rebalance
        )

    def coarse_selection_function(self, coarse):
        # top 50 by dollar volume
        selected = [x.symbol for x in sorted(coarse, key=lambda c: c.dollar_volume, reverse=True) if x.has_fundamental_data]
        return selected[:self.num_coarse]

    def fine_selection_function(self, fine):
        # pick top 20 stocks with positive PE
        sorted_by_pe = sorted(
            fine,
            key=lambda f: f.valuation_ratios.pe_ratio if f.valuation_ratios.pe_ratio > 0 else 9999
        )
        return [x.symbol for x in sorted_by_pe[:self.num_fine]]

    def on_securities_changed(self, changes):
        for security in changes.added_securities:
            # attach ATR indicator for each stock
            atr_indicator = self.atr(security.symbol, self.atr_period, MovingAverageType.SIMPLE, Resolution.DAILY)
            self.data[security.symbol] = atr_indicator

    def get_vix_adjustment(self):
        if not self.securities[self.vix].has_data:
            return 1.0
        
        vix_price = self.securities[self.vix].price
        if vix_price < 20:
            return 1.0   # normal risk
        elif vix_price < 30:
            return 0.5   # cut size in half
        else:
            return 0.25  # very defensive

    def rebalance(self):
        # liquidate everything first
        self.liquidate()
        
        selected = [s for s in self.data if self.securities.contains_key(s) and self.securities[s].has_data]
        
        # momentum filter (20d > 50d)
        final = []
        for symbol in selected:
            hist = self.history(symbol, 60, Resolution.DAILY)
            if len(hist) < 60: 
                continue
            closes = hist['close'].values
            ma20 = closes[-20:].mean()
            ma50 = closes[-50:].mean()
            if ma20 > ma50:
                final.append(symbol)
        
        # take up to 10 stocks
        final = final[:self.max_positions]
        if not final: 
            return
        
        # adjust weights by volatility regime
        vix_adj = self.get_vix_adjustment()
        weight = (1.0 / len(final)) * vix_adj
        
        for symbol in final:
            self.set_holdings(symbol, weight)

    def on_data(self, data):
        # ATR stop-loss check
        for kvp in self.portfolio:
            holding = kvp.value
            symbol = holding.symbol
            if not holding.invested: 
                continue
            if symbol not in self.data: 
                continue
            
            atr = self.data[symbol].current.value
            if atr == 0: 
                continue
            stop_price = holding.average_price - self.stop_loss_atr * atr
            if self.securities[symbol].price < stop_price:
                self.liquidate(symbol)