# Merged version of VolatilitySurfaceMomentum and MultiAssetStraddle strategies (Option 2: Parallel Execution)

from AlgorithmImports import *
from scipy.interpolate import SmoothBivariateSpline
from sklearn.decomposition import PCA
from datetime import timedelta
import numpy as np
from collections import deque

class IntegratedVolStraddle(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 1)
        self.SetEndDate(2019, 1, 1)
        self.SetCash(10_000_000)

        self.tickers = ["AAPL", "TSLA", "NVDA", "AMZN", "SPY"]
        self.under = {}
        self.option_symbol = {}
        self.positions = {}
        self.entryCost = {}
        self.entryExpiry = {}
        self.entrySpot = {}
        self.lastEntryTime = {}
        self.positionSource = {}  # track which strategy opened it

        # Strategy-specific parameters
        self.atr = {}
        self.macd = {}
        self.rsi = {}
        self.atrMultiplier = 1.5
        self.takeProfit = 0.4
        self.stopLoss = -0.05
        self.cooldown_days = 30
        self.max_trade_value = 50_000

        # IV surface momentum data
        self.spline = {}
        self.splineHist = {t: deque(maxlen=5) for t in self.tickers}
        self.hist = {f: {t: deque(maxlen=250) for t in self.tickers}
                     for f in ("level_change", "rolling_level", "skew_change", "smile_change")}
        self.q = {
            "level_change": 0.6448,
            "rolling_level": 0.6732,
            "smile_change": 0.5417,
            "skew_change": 0.6163
        }

        self.SetWarmUp(200)

        for t in self.tickers:
            eq = self.AddEquity(t, Resolution.Daily).Symbol
            op = self.AddOption(t, Resolution.Daily)
            op.SetFilter(lambda u: u.IncludeWeeklys().Strikes(-1, 1).Expiration(30, 45))
            self.under[t] = eq
            self.option_symbol[t] = op.Symbol
            self.atr[t] = self.ATR(eq, 14, Resolution.Daily)
            self.macd[t] = self.MACD(eq, 12, 26, 9, MovingAverageType.Wilders, Resolution.Daily)
            self.rsi[t] = self.RSI(eq, 14, MovingAverageType.Wilders, Resolution.Daily)

        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen("SPY", 10), self.FitSurfaces)

    def FitSurfaces(self):
        for t, sym in self.option_symbol.items():
            chain = self.OptionChain(sym)
            if not chain:
                continue
            c = [x for x in chain if x.ImpliedVolatility > 0 and x.UnderlyingLastPrice > 0]
            if len(c) < 15:
                continue
            try:
                K = [x.Strike for x in c]
                S = [x.UnderlyingLastPrice for x in c]
                M = [k / s - 1 for k, s in zip(K, S)]
                DTE = [(x.Expiry.date() - self.Time.date()).days for x in c]
                IV = [x.ImpliedVolatility for x in c]
                spl = SmoothBivariateSpline(M, DTE, IV, kx=3, ky=3, s=25)
                self.spline[t] = spl
                self.splineHist[t].append(spl)
            except Exception as e:
                self.Debug(f"{t} spline err {e}")

    def OnData(self, slice: Slice):
        if self.IsWarmingUp:
            return

        for t in self.tickers:
            if self.HasOpenStraddle(t):
                self._exit_checks(t)
                continue
            if self.lastEntryTime.get(t) and (self.Time - self.lastEntryTime[t]).days < self.cooldown_days:
                continue

            # MAS entry condition
            if self._atr_breakout(t) and self._macd_cross(t) and self._rsi_neutral(t):
                self._enter_straddle(t, slice, "MAS")

            # VSM entry condition
            if self._vsm_trigger(t):
                self._enter_straddle(t, slice, "VSM")

    def _atr_breakout(self, t):
        return self.atr[t].IsReady and (self.Securities[t].Close - self.Securities[t].Low) > self.atrMultiplier * self.atr[t].Current.Value

    def _macd_cross(self, t):
        return self.macd[t].IsReady and self.macd[t].Current.Value > self.macd[t].Signal.Current.Value

    def _rsi_neutral(self, t):
        return self.rsi[t].IsReady and 35 < self.rsi[t].Current.Value < 65

    def _vsm_trigger(self, t):
        if t not in self.spline or len(self.splineHist[t]) < 4:
            return False
        try:
            splines = list(self.splineHist[t])[-4:]
            m = np.linspace(-0.25, 0.25, 11)
            iv = [np.array([s.ev(x, 30) for x in m]) for s in splines]
            pca = PCA(3)
            pca.fit(iv[:3])
            pc1 = pca.transform([iv[1]])[0]
            pc2 = pca.transform([iv[2]])[0]
            pc3 = pca.transform([iv[3]])[0]
            d = pc2 - pc1
            lvl, skw, sml = d
            self.hist["level_change"][t].append(abs(lvl))
            self.hist["skew_change"][t].append(abs(skw))
            self.hist["smile_change"][t].append(abs(sml))
            rollLvl = np.mean(list(self.hist["level_change"][t])[-3:])
            self.hist["rolling_level"][t].append(rollLvl)
            qok = lambda f, val: abs(val) > np.quantile(self.hist[f][t], self.q[f])
            κ = np.mean(np.gradient(np.gradient(iv[-1], m), m) - np.gradient(np.gradient(iv[-2], m), m))
            return all([qok("level_change", lvl), qok("rolling_level", rollLvl), qok("skew_change", skw), qok("smile_change", sml)]) and abs(κ) > 0.01
        except:
            return False

    def _enter_straddle(self, t, slice, source):
        chain = slice.OptionChains.get(self.option_symbol[t])
        if not chain:
            return
        contracts = [c for c in chain if c.Expiry > self.Time]
        if not contracts:
            return
        expiry = sorted({c.Expiry for c in contracts})[0]
        spot = self.Securities[t].Price
        atm_strike = min(contracts, key=lambda c: abs(c.Strike - spot)).Strike
        calls = [c for c in contracts if c.Strike == atm_strike and c.Right == OptionRight.Call and c.AskPrice > 0 and c.Volume > 0 and c.Expiry == expiry]
        puts = [c for c in contracts if c.Strike == atm_strike and c.Right == OptionRight.Put and c.AskPrice > 0 and c.Volume > 0 and c.Expiry == expiry]
        if not calls or not puts:
            return
        call = sorted(calls, key=lambda c: (-c.Volume, c.AskPrice))[0]
        put = sorted(puts, key=lambda c: (-c.Volume, c.AskPrice))[0]
        mult = self.Securities[call.Symbol].SymbolProperties.ContractMultiplier
        price = (call.AskPrice + put.AskPrice) * mult
        if price == 0:
            return
        qty = int(self.max_trade_value / price)
        totalCost = qty * price
        if qty < 1 or totalCost > self.Portfolio.Cash * 0.9:
            return
        self.MarketOrder(call.Symbol, qty)
        self.MarketOrder(put.Symbol, qty)
        self.entryCost[t] = totalCost
        self.entryExpiry[t] = expiry
        self.entrySpot[t] = spot
        self.lastEntryTime[t] = self.Time
        self.positionSource[t] = source
        self.Debug(f"[{source}] Entered {t} straddle | Spot: {spot:.2f}, Qty: {qty}")

    def _exit_checks(self, t):
        positions = [h for h in self.Portfolio.Values if h.Invested and h.Symbol.SecurityType == SecurityType.Option and h.Symbol.Underlying.Value == t]
        if not positions:
            return
        unreal = sum(h.UnrealizedProfit for h in positions)
        cost = self.entryCost.get(t, 1)
        pl_pct = unreal / cost
        days_to_expiry = (self.entryExpiry[t].date() - self.Time.date()).days
        if pl_pct >= self.takeProfit or pl_pct <= self.stopLoss or days_to_expiry <= 1:
            self.Debug(f"[EXIT] {t} | PnL: {pl_pct:.2%}, DTE: {days_to_expiry}")
            self.Liquidate(t)

    def HasOpenStraddle(self, t):
        return any(h.Invested and h.Symbol.SecurityType == SecurityType.Option and h.Symbol.Underlying.Value == t for h in self.Portfolio.Values)