from AlgorithmImports import *

class TechMomentumWithHedges(QCAlgorithm):
    """
    High-growth tech momentum strategy with inverse correlation hedging
    - 60% high-growth tech (AMD, PLTR, NVDA, TSLA, etc.)
    - 40% negatively correlated hedges (gold, utilities, treasuries, inverse tech)
    - Goal: Capture tech upside while smoothing volatility for Sharpe 2.0+
    """

    def Initialize(self):
        self.SetStartDate(2024, 10, 1)
        self.SetEndDate(2025, 12, 31)
        self.SetCash(1000000)

        # HIGH GROWTH TECH - The Alpha Generators
        self.tech_stocks = [
            "NVDA",   # AI/GPU leader
            "AMD",    # CPU/GPU competitor
            "PLTR",   # AI software/data
            "TSLA",   # EV/autonomy
            "MSFT",   # Cloud/AI
            "GOOGL",  # AI/search
            "META",   # Social/AI
            "AVGO",   # Semiconductors
            "ARM",    # Chip design
            "SNOW",   # Cloud data
        ]

        # INVERSE/HEDGE ASSETS - Volatility Dampeners
        self.hedge_assets = {
            "GLD": 0.12,    # Gold (inverse to tech risk-on)
            "TLT": 0.12,    # Long-term treasuries
            "XLU": 0.08,    # Utilities sector (defensive)
            "VNQ": 0.04,    # Real estate
            "SQQQ": 0.14,   # 3x Inverse Nasdaq (direct tech hedge)
        }

        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.SetBenchmark("SPY")

        # Add all securities
        self.symbols = {}
        for ticker in self.tech_stocks + list(self.hedge_assets.keys()):
            symbol = self.AddEquity(ticker, Resolution.Daily).Symbol
            self.symbols[ticker] = symbol

        # Momentum indicators for tech stocks
        self.indicators = {}
        for ticker in self.tech_stocks:
            symbol = self.symbols[ticker]
            self.indicators[ticker] = {
                'rsi': self.RSI(symbol, 14),
                'macd': self.MACD(symbol, 12, 26, 9),
                'sma20': self.SMA(symbol, 20),
                'sma50': self.SMA(symbol, 50),
                'mom': self.MOMP(symbol, 21),  # 21-day momentum
                'std': self.STD(symbol, 20)
            }

        # Track performance
        self.recent_returns = {}

        # Rebalance weekly (sweet spot between daily noise and monthly lag)
        self.Schedule.On(
            self.DateRules.WeekStart(self.spy),
            self.TimeRules.AfterMarketOpen(self.spy, 30),
            self.Rebalance
        )

        # Risk management daily
        self.Schedule.On(
            self.DateRules.EveryDay(self.spy),
            self.TimeRules.AfterMarketOpen(self.spy, 60),
            self.RiskManagement
        )

        # ⚡ Start immediately, no warm-up delay
        self.SetWarmUp(0)

        # Force first rebalance immediately on start date
        self.Schedule.On(
            self.DateRules.Today,
            self.TimeRules.AfterMarketOpen(self.spy, 1),
            self.Rebalance
        )

        self.entry_prices = {}

    def CalculateMomentumScore(self, ticker):
        """
        Score tech stocks on momentum strength (0-100)
        Higher score = stronger momentum = larger position
        """
        if not self.indicators[ticker]['macd'].IsReady:
            return 0

        symbol = self.symbols[ticker]
        price = self.Securities[symbol].Price

        ind = self.indicators[ticker]
        rsi = ind['rsi'].Current.Value
        macd = ind['macd'].Current.Value
        macd_signal = ind['macd'].Signal.Current.Value
        sma20 = ind['sma20'].Current.Value
        sma50 = ind['sma50'].Current.Value
        momentum = ind['mom'].Current.Value

        score = 0

        # FACTOR 1: Trend strength (40 points)
        if price > sma20 and price > sma50:
            score += 20
        if sma20 > sma50:  # Short-term above long-term
            score += 20

        # FACTOR 2: MACD momentum (30 points)
        if macd > macd_signal and macd > 0:
            score += 30
        elif macd > macd_signal:
            score += 15

        # FACTOR 3: RSI sweet spot (20 points)
        if 45 < rsi < 70:  # Strong but not overbought
            score += 20
        elif 40 < rsi < 75:
            score += 10

        # FACTOR 4: Recent momentum (10 points)
        if momentum > 0:
            score += 10

        return score

    def Rebalance(self):
        """Weekly rebalancing: momentum-weighted tech + fixed hedges"""
        if self.IsWarmingUp:
            return

        # Calculate momentum scores for all tech stocks
        tech_scores = {}
        for ticker in self.tech_stocks:
            score = self.CalculateMomentumScore(ticker)
            tech_scores[ticker] = score

        # Rank and select top 6 tech stocks
        sorted_tech = sorted(tech_scores.items(), key=lambda x: x[1], reverse=True)
        top_tech = [ticker for ticker, score in sorted_tech[:6] if score > 50]

        # Build target portfolio
        target_weights = {}

        # TECH ALLOCATION (60% total, momentum-weighted)
        if len(top_tech) > 0:
            total_score = sum([tech_scores[t] for t in top_tech])

            for ticker in top_tech:
                if total_score > 0:
                    # Weight proportional to momentum score
                    raw_weight = (tech_scores[ticker] / total_score) * 0.60

                    # Apply volatility adjustment
                    volatility = self.indicators[ticker]['std'].Current.Value
                    price = self.Securities[self.symbols[ticker]].Price

                    if volatility > 0 and price > 0:
                        vol_pct = volatility / price
                        # Reduce weight for high-vol stocks
                        vol_adjustment = max(0.6, 1 - vol_pct * 5)
                        adjusted_weight = raw_weight * vol_adjustment
                    else:
                        adjusted_weight = raw_weight

                    # Cap individual positions at 15%
                    target_weights[ticker] = min(adjusted_weight, 0.15)

        # Normalize tech weights to exactly 60%
        total_tech_weight = sum([target_weights.get(t, 0) for t in self.tech_stocks])
        if total_tech_weight > 0:
            tech_scale = 0.60 / total_tech_weight
            for ticker in self.tech_stocks:
                if ticker in target_weights:
                    target_weights[ticker] *= tech_scale

        # DYNAMIC HEDGE ALLOCATION (adapt to market volatility)
        tech_vols = []
        for ticker in top_tech:
            if ticker in self.indicators and self.indicators[ticker]['std'].IsReady:
                vol = self.indicators[ticker]['std'].Current.Value
                price = self.Securities[self.symbols[ticker]].Price
                if price > 0:
                    tech_vols.append(vol / price)

        avg_tech_vol = sum(tech_vols) / len(tech_vols) if tech_vols else 0.02

        # Adaptive hedge ratio: higher vol = more hedging
        if avg_tech_vol > 0.035:  # High volatility regime
            hedge_ratio = 0.65  # 65% hedges, 35% tech
            tech_ratio = 0.35
        elif avg_tech_vol > 0.025:  # Medium volatility
            hedge_ratio = 0.50  # 50/50 split
            tech_ratio = 0.50
        else:  # Low volatility regime
            hedge_ratio = 0.40  # 40% hedges, 60% tech
            tech_ratio = 0.60

        # Rescale tech positions
        if total_tech_weight > 0:
            tech_scale = tech_ratio / total_tech_weight
            for ticker in self.tech_stocks:
                if ticker in target_weights:
                    target_weights[ticker] *= tech_scale

        # Scale hedge allocations
        hedge_scale = hedge_ratio / sum(self.hedge_assets.values())
        for ticker, base_weight in self.hedge_assets.items():
            target_weights[ticker] = base_weight * hedge_scale

        self.Debug(f"Volatility regime: {avg_tech_vol*100:.2f}% | Hedges: {hedge_ratio*100:.0f}%")

        # Execute trades
        for ticker in self.tech_stocks + list(self.hedge_assets.keys()):
            symbol = self.symbols[ticker]
            target_weight = target_weights.get(ticker, 0)

            current_weight = 0
            if self.Portfolio[symbol].Invested:
                current_weight = self.Portfolio[symbol].HoldingsValue / self.Portfolio.TotalPortfolioValue

            # Trade if deviation > 2%
            if abs(current_weight - target_weight) > 0.02:
                self.SetHoldings(symbol, target_weight)

                # Track entry for risk management
                if target_weight > current_weight and ticker in self.tech_stocks:
                    self.entry_prices[ticker] = self.Securities[symbol].Price

        # Log allocation
        total_tech = sum([target_weights.get(t, 0) for t in self.tech_stocks])
        total_hedge = sum([target_weights.get(t, 0) for t in self.hedge_assets.keys()])
        self.Debug(f"Portfolio: {total_tech*100:.1f}% tech, {total_hedge*100:.1f}% hedges")

    def RiskManagement(self):
        """Daily risk checks: profit-taking and stop-losses on tech positions"""
        if self.IsWarmingUp:
            return

        for ticker in self.tech_stocks:
            symbol = self.symbols[ticker]

            if not self.Portfolio[symbol].Invested:
                continue

            current_price = self.Securities[symbol].Price
            entry_price = self.entry_prices.get(ticker)

            if not entry_price or entry_price == 0:
                continue

            pnl_pct = (current_price - entry_price) / entry_price

            # Take partial profits on big winners
            if pnl_pct > 0.20:  # Up 20%
                quantity = self.Portfolio[symbol].Quantity
                self.MarketOrder(symbol, -int(quantity * 0.3))
                self.Debug(f"Taking 30% profit on {ticker} at +{pnl_pct*100:.1f}%")

            elif pnl_pct > 0.35:  # Up 35%
                quantity = self.Portfolio[symbol].Quantity
                self.MarketOrder(symbol, -int(quantity * 0.4))
                self.Debug(f"Taking 40% profit on {ticker} at +{pnl_pct*100:.1f}%")

            # Stop-loss to limit damage
            elif pnl_pct < -0.08:  # Down 8%
                self.Liquidate(symbol)
                self.Debug(f"Stop-loss on {ticker} at {pnl_pct*100:.1f}%")
                if ticker in self.entry_prices:
                    del self.entry_prices[ticker]

            # Trailing stop on winners
            elif pnl_pct > 0.12:  # After +12% gain
                # 6% trailing stop from peak
                peak_price = entry_price * 1.12
                if current_price < peak_price * 0.94:
                    self.Liquidate(symbol)
                    self.Debug(f"Trailing stop on {ticker}")
                    if ticker in self.entry_prices:
                        del self.entry_prices[ticker]

    def OnData(self, data):
        pass

    def OnEndOfAlgorithm(self):
        self.Debug(f"Final Portfolio Value: ${self.Portfolio.TotalPortfolioValue:,.2f}")
        self.Debug(f"Total Return: {(self.Portfolio.TotalPortfolioValue/1000000 - 1)*100:.2f}%")