from AlgorithmImports import *
import numpy as np
from datetime import timedelta


class GoldVolatilityArbitrageStrategy(QCAlgorithm):
    """
    Institutional gold-focused strategy:
    1. Gold mean reversion on extreme moves
    2. Gold miners arbitrage (GDX/GDXJ vs GLD spread)
    3. Real rates regime trading (Gold vs TLT)
    4. Cross-asset correlations (Gold/USD, Gold/Equity risk)
    """

    def Initialize(self):
        self.SetStartDate(2024, 1, 1)
        self.SetCash(10_000_000)

        # Core gold instruments - DAILY resolution
        self.gld = self.AddEquity("GLD", Resolution.Daily).Symbol
        self.gdx = self.AddEquity("GDX", Resolution.Daily).Symbol
        self.gdxj = self.AddEquity("GDXJ", Resolution.Daily).Symbol
        self.slv = self.AddEquity("SLV", Resolution.Daily).Symbol

        # Macro instruments
        self.tlt = self.AddEquity("TLT", Resolution.Daily).Symbol
        self.uup = self.AddEquity("UUP", Resolution.Daily).Symbol
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.tip = self.AddEquity("TIP", Resolution.Daily).Symbol

        # Volatility proxy
        self.vxx = self.AddEquity("VXX", Resolution.Daily).Symbol

        # Data storage
        self.data = {}
        for symbol in [self.gld, self.gdx, self.gdxj, self.slv, self.tlt, self.uup, self.spy, self.tip, self.vxx]:
            self.data[symbol] = SymbolData(symbol, self, lookback=120)

        # Strategy state
        self.gold_zscore = 0.0
        self.miner_spread_zscore = 0.0
        self.real_rate_regime = "NEUTRAL"
        self.usd_regime = "NEUTRAL"
        self.crash_protection_active = False

        # Risk parameters
        self.max_gross_leverage = 2.5
        self.volatility_target = 0.15
        self.max_drawdown = 0.10
        self.peak_equity = self.Portfolio.TotalPortfolioValue

        # Reversion parameters
        self.gold_reversion_threshold = 2.0
        self.gold_exit_threshold = 0.5
        self.miner_arb_threshold = 1.5

        # Warmup
        self.SetWarmUp(60, Resolution.Daily)

        # Scheduling (QC uses PascalCase DateRules/TimeRules/Schedule)
        self.Schedule.On(
            self.DateRules.EveryDay(self.gld),
            self.TimeRules.AfterMarketOpen(self.gld, 30),
            self.AnalyzeMarketRegime
        )

        self.Schedule.On(
            self.DateRules.EveryDay(self.gld),
            self.TimeRules.AfterMarketOpen(self.gld, 60),
            self.ExecuteGoldMeanReversion
        )

        self.Schedule.On(
            self.DateRules.EveryDay(self.gld),
            self.TimeRules.AfterMarketOpen(self.gld, 120),
            self.ExecuteMinerArbitrage
        )

        self.Schedule.On(
            self.DateRules.EveryDay(self.gld),
            self.TimeRules.BeforeMarketClose(self.gld, 30),
            self.RiskManagement
        )

        self.Debug("Gold strategy initialized (WarmUp: 60 days).")

    def AnalyzeMarketRegime(self):
        if self.IsWarmingUp:
            return

        needed = [self.gld, self.tlt, self.uup, self.spy, self.tip]
        if not all(self.data[s].IsReady() for s in needed):
            return

        gld_prices = self.data[self.gld].GetPrices()
        if gld_prices is None or len(gld_prices) < 30:
            return

        # Z-score (shorter lookback for more signals)
        lookback = min(40, len(gld_prices))
        window = gld_prices[-lookback:]
        mu = float(np.mean(window))
        sd = float(np.std(window))
        self.gold_zscore = (float(gld_prices[-1]) - mu) / sd if sd > 1e-12 else 0.0

        # Real rates proxy: TLT/TIP ratio trend (align lengths)
        tip_prices = self.data[self.tip].GetPrices()
        tlt_prices = self.data[self.tlt].GetPrices()
        if tip_prices is not None and tlt_prices is not None:
            n = min(len(tip_prices), len(tlt_prices))
            if n >= 20:
                tip = tip_prices[-20:]
                tlt = tlt_prices[-20:]
                real_rate_proxy = tlt / (tip + 1e-10)
                real_rate_trend = float(real_rate_proxy[-1] / real_rate_proxy[0] - 1.0)

                if real_rate_trend > 0.02:
                    self.real_rate_regime = "RISING_RATES"
                elif real_rate_trend < -0.02:
                    self.real_rate_regime = "FALLING_RATES"
                else:
                    self.real_rate_regime = "NEUTRAL"

        # USD regime via SMA(20) vs SMA(50)
        uup_prices = self.data[self.uup].GetPrices()
        if uup_prices is not None and len(uup_prices) >= 50:
            sma20 = float(np.mean(uup_prices[-20:]))
            sma50 = float(np.mean(uup_prices[-50:]))

            if sma20 > sma50 * 1.02:
                self.usd_regime = "USD_BULL"
            elif sma20 < sma50 * 0.98:
                self.usd_regime = "USD_BEAR"
            else:
                self.usd_regime = "NEUTRAL"

        # Crash detection from SPY 5-day return
        spy_rets = self.data[self.spy].GetReturns()
        if spy_rets is not None and len(spy_rets) >= 5:
            recent = float(np.sum(spy_rets[-5:]))
            if recent < -0.05:
                self.crash_protection_active = True
            elif recent > 0.03:
                self.crash_protection_active = False

    def ExecuteGoldMeanReversion(self):
        if self.IsWarmingUp:
            return
        if not self.data[self.gld].IsReady():
            return

        gld_prices = self.data[self.gld].GetPrices()
        if gld_prices is None or len(gld_prices) < 20:
            return

        gld_rets = self.data[self.gld].GetReturns()
        if gld_rets is None or len(gld_rets) < 20:
            return

        # Short-term momentum
        short_term_return = float(gld_prices[-1] / gld_prices[-5] - 1.0) if len(gld_prices) >= 5 else 0.0

        # Volatility scaling
        ann_vol = float(np.std(gld_rets[-20:]) * np.sqrt(252))
        ann_vol = max(ann_vol, 1e-6)
        vol_scalar = float(np.clip(self.volatility_target / ann_vol, 0.5, 2.0))

        tpv = self.Portfolio.TotalPortfolioValue
        current_weight = float(self.Portfolio[self.gld].HoldingsValue / tpv) if tpv > 0 else 0.0

        # Extreme plummet: long
        if self.gold_zscore < -self.gold_reversion_threshold:
            target = 0.40 * vol_scalar
            if self.real_rate_regime == "FALLING_RATES":
                target *= 1.3
            if self.usd_regime == "USD_BEAR":
                target *= 1.2
            if self.crash_protection_active:
                target *= 1.5
            target = float(min(target, 0.80))

            if abs(current_weight - target) > 0.05:
                self.SetHoldings(self.gld, target)

        # Extreme spike: short
        elif self.gold_zscore > self.gold_reversion_threshold:
            target = -0.30 * vol_scalar
            if self.real_rate_regime == "RISING_RATES":
                target *= 1.3
            if self.usd_regime == "USD_BULL":
                target *= 1.2
            target = float(max(target, -0.50))

            if abs(current_weight - target) > 0.05:
                self.SetHoldings(self.gld, target)

        # Exit on reversion: reduce exposure (target closer to 0)
        elif abs(self.gold_zscore) < self.gold_exit_threshold and abs(current_weight) > 0.10:
            target = float(current_weight * 0.70)  # reduce 30%
            self.SetHoldings(self.gld, target)

        # Baseline small exposure when quiet
        elif abs(current_weight) < 0.05 and abs(self.gold_zscore) < 1.0:
            self.SetHoldings(self.gld, 0.15 if short_term_return > 0 else -0.10)

    def ExecuteMinerArbitrage(self):
        if self.IsWarmingUp:
            return
        if not all(self.data[s].IsReady() for s in [self.gld, self.gdx, self.gdxj]):
            return

        gld = self.data[self.gld].GetPrices()
        gdx = self.data[self.gdx].GetPrices()
        gdxj = self.data[self.gdxj].GetPrices()
        if gld is None or gdx is None or gdxj is None:
            return

        # Align lengths to avoid numpy shape mismatch
        n = min(len(gld), len(gdx), len(gdxj))
        if n < 40:
            return
        gld = gld[-n:]
        gdx = gdx[-n:]
        gdxj = gdxj[-n:]

        ratio = gdx / (gld + 1e-10)

        lookback = min(40, len(ratio))
        window = ratio[-lookback:]
        mu = float(np.mean(window))
        sd = float(np.std(window))
        self.miner_spread_zscore = (float(ratio[-1]) - mu) / (sd + 1e-10)

        tpv = self.Portfolio.TotalPortfolioValue
        current_gdx_w = float(self.Portfolio[self.gdx].HoldingsValue / tpv) if tpv > 0 else 0.0

        # Miners cheap: long miners (optionally reduce gold elsewhere)
        if self.miner_spread_zscore < -self.miner_arb_threshold:
            self.SetHoldings(self.gdx, 0.25)
            self.SetHoldings(self.gdxj, 0.15)

        # Miners expensive: short miners
        elif self.miner_spread_zscore > self.miner_arb_threshold:
            self.SetHoldings(self.gdx, -0.15)
            if self.Portfolio[self.gdxj].Invested:
                self.Liquidate(self.gdxj)

        # Close arb if converged (and not in strong main gold signal)
        elif abs(self.miner_spread_zscore) < 0.5:
            if abs(current_gdx_w) > 0.05 and abs(self.gold_zscore) <= 2.0:
                self.Liquidate(self.gdx)
                self.Liquidate(self.gdxj)

    def RiskManagement(self):
        if self.IsWarmingUp:
            return

        tpv = self.Portfolio.TotalPortfolioValue
        if tpv <= 0:
            return

        # Drawdown tracking
        if tpv > self.peak_equity:
            self.peak_equity = tpv
        dd = float((self.peak_equity - tpv) / self.peak_equity) if self.peak_equity > 0 else 0.0

        # Drawdown protection: scale down these exposures
        if dd > self.max_drawdown:
            scale = 0.5
            for symbol in [self.gld, self.gdx, self.gdxj, self.slv]:
                if self.Portfolio[symbol].Invested:
                    w = float(self.Portfolio[symbol].HoldingsValue / tpv)
                    self.SetHoldings(symbol, w * scale)

        # Gross leverage estimate (manual): sum(abs(weights))
        gross = 0.0
        for sym in [self.gld, self.gdx, self.gdxj, self.slv, self.tlt, self.uup, self.spy, self.tip, self.vxx]:
            if self.Portfolio[sym].Invested:
                gross += abs(float(self.Portfolio[sym].HoldingsValue / tpv))

        if gross > self.max_gross_leverage:
            scale = float(self.max_gross_leverage / gross)
            for sym in [self.gld, self.gdx, self.gdxj, self.slv, self.tlt, self.uup, self.spy, self.tip, self.vxx]:
                if self.Portfolio[sym].Invested:
                    w = float(self.Portfolio[sym].HoldingsValue / tpv)
                    self.SetHoldings(sym, w * scale)

        # Correlation risk (GLD vs GDX)
        if self.Portfolio[self.gld].Invested and self.Portfolio[self.gdx].Invested:
            gld_r = self.data[self.gld].GetReturns()
            gdx_r = self.data[self.gdx].GetReturns()
            if gld_r is not None and gdx_r is not None:
                n = min(len(gld_r), len(gdx_r))
                if n >= 20:
                    corr = float(np.corrcoef(gld_r[-n:], gdx_r[-n:])[0, 1])
                    if abs(corr) > 0.90:
                        gld_val = abs(float(self.Portfolio[self.gld].HoldingsValue))
                        gdx_val = abs(float(self.Portfolio[self.gdx].HoldingsValue))
                        if gdx_val < gld_val:
                            self.Liquidate(self.gdx)


class SymbolData:
    def __init__(self, symbol: Symbol, algorithm: QCAlgorithm, lookback: int = 120):
        self.Symbol = symbol
        self.algorithm = algorithm
        self.lookback = lookback

        self.prices = RollingWindow[float](lookback)
        self.volumes = RollingWindow[float](lookback)

        self.consolidator = TradeBarConsolidator(timedelta(days=1))
        self.consolidator.DataConsolidated += self.OnDataConsolidated
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)

    def OnDataConsolidated(self, sender, bar: TradeBar):
        self.prices.Add(float(bar.Close))
        self.volumes.Add(float(bar.Volume))

    def IsReady(self) -> bool:
        return self.prices.IsReady

    def GetPrices(self):
        if not self.IsReady():
            return None
        n = min(self.lookback, self.prices.Count)
        return np.array([self.prices[i] for i in range(n)], dtype=float)

    def GetReturns(self):
        prices = self.GetPrices()
        if prices is None or len(prices) < 2:
            return None
        return np.diff(np.log(prices))

    def Dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.consolidator)