from AlgorithmImports import *

class TechMomentumWithHedges(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2021, 1, 1)
        self.SetCash(100_000)
        
        self.tech_stocks = ["NVDA", "AMD", "PLTR", "TSLA", "MSFT", "GOOGL", "META", "AVGO", "ARM", "SNOW"]
        self.hedge_assets = {"GLD": 0.12, "TLT": 0.12, "XLU": 0.08, "VNQ": 0.04, "SQQQ": 0.14}
        
        # Benchmark and Relative Strength Indicator
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.spy_mom = self.MOMP(self.spy, 21)
        self.SetBenchmark(self.spy)
        
        self.symbols = {}
        for ticker in self.tech_stocks + list(self.hedge_assets.keys()):
            self.symbols[ticker] = self.AddEquity(ticker, Resolution.Daily).Symbol
        
        self.indicators = {}
        for ticker in self.tech_stocks:
            s = self.symbols[ticker]
            self.indicators[ticker] = {
                'rsi': self.RSI(s, 14),
                'macd': self.MACD(s, 12, 26, 9),
                'sma20': self.SMA(s, 20),
                'sma50': self.SMA(s, 50),
                'mom': self.MOMP(s, 21),
                'std': self.STD(s, 20)
            }
        
        self.Schedule.On(self.DateRules.WeekStart(self.spy), self.TimeRules.AfterMarketOpen(self.spy, 30), self.Rebalance)
        self.Schedule.On(self.DateRules.EveryDay(self.spy), self.TimeRules.AfterMarketOpen(self.spy, 60), self.RiskManagement)
        
        self.SetWarmUp(50)
        self.set_brokerage_model(BrokerageName.ALPACA) 

    def CalculateMomentumScore(self, ticker):
        if not self.indicators[ticker]['macd'].IsReady or not self.spy_mom.IsReady:
            return 0
        
        price = self.Securities[self.symbols[ticker]].Price
        ind = self.indicators[ticker]
        score = 0
        
        # Trend (40pts)
        if price > ind['sma20'].Current.Value and price > ind['sma50'].Current.Value: score += 20
        if ind['sma20'].Current.Value > ind['sma50'].Current.Value: score += 20
        
        # MACD (20pts)
        if ind['macd'].Current.Value > ind['macd'].Signal.Current.Value: score += 20
        
        # RSI Sweet Spot (20pts)
        rsi = ind['rsi'].Current.Value
        if 45 < rsi < 70: score += 20
        
        # REFINED: Relative Strength (20pts)
        # Bonus if outperforming SPY over the last month
        if ind['mom'].Current.Value > self.spy_mom.Current.Value:
            score += 20
            
        return score

    def Rebalance(self):
        if self.IsWarmingUp: return
        
        # 1. DETERMINE REGIME
        tech_vols = [self.indicators[t]['std'].Current.Value / self.Securities[self.symbols[t]].Price 
                     for t in self.tech_stocks if self.indicators[t]['std'].IsReady]
        avg_vol = sum(tech_vols) / len(tech_vols) if tech_vols else 0.02
        
        if avg_vol > 0.035:   tech_ratio, hedge_ratio = 0.35, 0.65
        elif avg_vol > 0.025: tech_ratio, hedge_ratio = 0.50, 0.50
        else:                 tech_ratio, hedge_ratio = 0.60, 0.40

        # 2. SELECT TECH TARGETS
        tech_scores = {t: self.CalculateMomentumScore(t) for t in self.tech_stocks}
        top_tech = [t for t, s in sorted(tech_scores.items(), key=lambda x: x[1], reverse=True)[:6] if s > 50]
        
        target_weights = {}

        # 3. ASSIGN TECH WEIGHTS (Mathematically locked to tech_ratio)
        if top_tech:
            total_score = sum(tech_scores[t] for t in top_tech)
            for t in top_tech:
                # Weighted by score, capped at 15% of total portfolio
                raw_weight = (tech_scores[t] / total_score) * tech_ratio
                target_weights[t] = min(raw_weight, 0.15)
            
            # Re-normalize tech to ensure they sum exactly to tech_ratio
            current_tech_sum = sum(target_weights.values())
            scale = tech_ratio / current_tech_sum
            for t in top_tech: target_weights[t] *= scale

        # 4. ASSIGN HEDGE WEIGHTS (Locked to hedge_ratio)
        base_hedge_sum = sum(self.hedge_assets.values())
        for ticker, base_w in self.hedge_assets.items():
            target_weights[ticker] = (base_w / base_hedge_sum) * hedge_ratio

        # 5. EXECUTION
        for ticker in self.tech_stocks + list(self.hedge_assets.keys()):
            symbol = self.symbols[ticker]
            weight = target_weights.get(ticker, 0)
            if abs(self.Portfolio[symbol].HoldingsValue / self.Portfolio.TotalPortfolioValue - weight) > 0.02:
                self.SetHoldings(symbol, weight)

    def RiskManagement(self):
        if self.IsWarmingUp: return
        
        for ticker in self.tech_stocks:
            symbol = self.symbols[ticker]
            if not self.Portfolio[symbol].Invested: continue
            
            # LOGIC CORRECTION: Use AveragePrice (Cost Basis)
            cost_basis = self.Portfolio[symbol].AveragePrice
            current_price = self.Securities[symbol].Price
            pnl_pct = (current_price - cost_basis) / cost_basis
            
            # Profit Taking (Partial)
            if pnl_pct > 0.30:
                self.MarketOrder(symbol, -int(self.Portfolio[symbol].Quantity * 0.5))
                self.Debug(f"Harvested 50% profit on {ticker}")
            
            # Stop Loss
            elif pnl_pct < -0.08:
                self.Liquidate(symbol)
                self.Debug(f"Stop-loss triggered for {ticker}")