from AlgorithmImports import *


class ExecutionGuard:
    """
    Avoid tiny trades and force MarketOnOpen to prevent QC daily-market-order warnings.
    """
    def __init__(self, minOrderValue: float = 2000, minWeightChange: float = 0.01) -> None:
        self.minOrderValue = minOrderValue
        self.minWeightChange = minWeightChange

    def ApplyTargets(self, algo: QCAlgorithm, targets: dict) -> None:
        """
        targets: dict[Symbol] -> target portfolio weight
        """
        total = float(algo.Portfolio.TotalPortfolioValue)
        if total <= 0:
            return

        # Liquidate anything not in targets
        keep = set(targets.keys())
        for kvp in algo.Portfolio:
            sym = kvp.Key
            if kvp.Value.Invested and sym not in keep:
                algo.Liquidate(sym)

        # Place changes with guards
        for sym, w in targets.items():
            if not algo.Securities.ContainsKey(sym):
                continue
            sec = algo.Securities[sym]
            if not sec.IsTradable or sec.Price is None or sec.Price <= 0:
                continue

            curWeight = float(algo.Portfolio[sym].HoldingsValue) / total if algo.Portfolio[sym].Invested else 0.0
            if abs(w - curWeight) < self.minWeightChange:
                continue

            targetValue = total * w
            if abs(targetValue - algo.Portfolio[sym].HoldingsValue) < self.minOrderValue:
                continue

            # Use MarketOnOpen explicitly for daily data
            algo.SetHoldings(sym, w, tag="MOO rebal")

# region imports
from AlgorithmImports import *
# endregion

import math
from execution import ExecutionGuard


class RobustCoreDefensiveOverlay(QCAlgorithm):
    """
    Core: SPY
    Defensive assets: SHY + GLD (stable crisis ballast)
    Optional tilt: XLE + ITA only in high-stress (rare)
    Regime:
      - Risk-on: SPY above 200SMA -> hold mostly SPY, vol-targeted
      - Risk-off: SPY below 200SMA -> hold defensive basket
    Rebalance monthly at market open using MOO behavior
    """

    def Initialize(self) -> None:
        self.SetStartDate(2006, 1, 1)
        self.SetCash(100000)

        self.UniverseSettings.Resolution = Resolution.DAILY
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.ADJUSTED

        # Assets
        self.spy = self.AddEquity("SPY", Resolution.DAILY).Symbol
        self.shy = self.AddEquity("SHY", Resolution.DAILY).Symbol   # 1-3y Treasuries
        self.gld = self.AddEquity("GLD", Resolution.DAILY).Symbol   # Gold
        self.xle = self.AddEquity("XLE", Resolution.DAILY).Symbol   # Energy
        self.ita = self.AddEquity("ITA", Resolution.DAILY).Symbol   # Defense

        # Indicators
        self.spySma200 = self.SMA(self.spy, 200, Resolution.DAILY)

        # Risk + sizing
        self.targetAnnVolRiskOn = 0.12
        self.targetAnnVolRiskOff = 0.06
        self.volLookback = 20

        # Stress tilt trigger (optional, rare)
        self.energySpikeLookback = 20
        self.energySpikeThreshold = 0.10  # +10% over 20d
        self.volHighThreshold = 0.020     # daily vol

        # Weights (base, before vol targeting)
        self.baseRiskOnSpy = 1.00
        self.baseRiskOff = {self.shy: 0.70, self.gld: 0.30}

        # In stress-on within risk-on, add small tilt
        self.stressTilt = {self.xle: 0.07, self.ita: 0.07}  # taken out of SPY

        # Execution guard to prevent tiny trades
        self.exec = ExecutionGuard(minOrderValue=2500, minWeightChange=0.01)

        # Monthly rebalance flag
        self.rebalanceFlag = True
        self.lastRebalanceMonth = None
        self.Schedule.On(
            self.DateRules.MonthStart("SPY"),
            self.TimeRules.AfterMarketOpen("SPY", 1),
            self.SetRebalanceFlag
        )

        self.SetWarmUp(260, Resolution.DAILY)

    def SetRebalanceFlag(self) -> None:
        self.rebalanceFlag = True

    def OnData(self, data: Slice) -> None:
        if self.IsWarmingUp:
            return

        if not self.rebalanceFlag:
            return

        monthKey = (self.Time.year, self.Time.month)
        if self.lastRebalanceMonth == monthKey:
            return
        self.lastRebalanceMonth = monthKey
        self.rebalanceFlag = False

        # Decide regime
        riskOn = self.spySma200.IsReady and self.Securities[self.spy].Price > self.spySma200.Current.Value

        if not riskOn:
            # Risk-off: defensive basket, vol-targeted
            targets = self.VolTarget(self.baseRiskOff, self.targetAnnVolRiskOff)
            self.exec.ApplyTargets(self, targets)
            return

        # Risk-on: mostly SPY, possibly small stress tilt
        targets = {self.spy: self.baseRiskOnSpy}

        if self.IsHighStress():
            # Take tilt weights from SPY
            take = sum(self.stressTilt.values())
            targets[self.spy] = max(0.0, 1.0 - take)
            for s, w in self.stressTilt.items():
                targets[s] = w

        # Vol target the whole portfolio
        targets = self.VolTarget(targets, self.targetAnnVolRiskOn)
        self.exec.ApplyTargets(self, targets)

    def IsHighStress(self) -> bool:
        """
        High stress proxy: oil/energy spike OR elevated SPY realized vol.
        """
        needed = max(self.energySpikeLookback + 2, self.volLookback + 2)
        hist = self.History([self.spy, self.xle], needed, Resolution.DAILY)
        if hist.empty:
            return False

        try:
            spyC = hist.loc[self.spy]["close"].dropna()
            xleC = hist.loc[self.xle]["close"].dropna()
        except Exception:
            return False

        if len(spyC) < needed or len(xleC) < needed:
            return False

        # SPY realized vol
        spyRet = spyC.pct_change().dropna()
        if len(spyRet) < self.volLookback:
            return False
        vol = float(spyRet.iloc[-self.volLookback:].std())
        volHigh = vol >= self.volHighThreshold

        # Energy spike
        xle20 = (float(xleC.iloc[-1]) / float(xleC.iloc[-1 - self.energySpikeLookback])) - 1.0
        energySpike = xle20 >= self.energySpikeThreshold

        return volHigh or energySpike

    def VolTarget(self, targets: dict, targetAnnVol: float) -> dict:
        """
        Scale weights to match target annual vol using diagonal covariance approximation:
        portDailyVol ~= sqrt(sum(w^2 * vol_i^2))
        """
        syms = list(targets.keys())
        if not syms:
            return {}

        hist = self.History(syms, self.volLookback + 2, Resolution.DAILY)
        if hist.empty:
            return targets

        vols = {}
        for s in syms:
            try:
                c = hist.loc[s]["close"].dropna()
                r = c.pct_change().dropna()
                if len(r) < self.volLookback:
                    continue
                vols[s] = float(r.iloc[-self.volLookback:].std())
            except Exception:
                continue

        # If missing vols, just return unscaled
        if len(vols) < max(1, len(syms) // 2):
            return targets

        portDailyVol = 0.0
        for s, w in targets.items():
            v = vols.get(s, None)
            if v is None or v <= 0:
                continue
            portDailyVol += (w * v) ** 2
        portDailyVol = math.sqrt(portDailyVol)

        if portDailyVol <= 0:
            return targets

        targetDaily = targetAnnVol / math.sqrt(252.0)
        scale = min(1.0, targetDaily / portDailyVol)

        scaled = {s: w * scale for s, w in targets.items()}

        # Normalize to <= 1 gross (keep conservative)
        ssum = sum(scaled.values())
        if ssum > 1.0 and ssum > 0:
            for k in list(scaled.keys()):
                scaled[k] /= ssum

        return scaled