from AlgorithmImports import *
from datetime import timedelta

class MultiAssetStraddle(QCAlgorithm):

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

        self.tickers = ["AAPL", "MSFT", "GOOG", "AMZN", "META", "TSLA", "NVDA"]

        self.atr = {}
        self.macd = {}
        self.rsi = {}
        self.option_symbol = {}
        self.entryCost = {}
        self.entryExpiry = {}
        self.entrySpot = {}
        self.lastEntryTime = {}

        self.atrMultiplier = 1.5
        self.takeProfit = 0.4
        self.stopLoss = -0.05
        self.cooldown_days = 30
        self.max_trade_value = 500_000

        self.SetWarmup(200)

        for ticker in self.tickers:
            equity = self.AddEquity(ticker, Resolution.Daily).Symbol
            option = self.AddOption(ticker, Resolution.Daily)
            option.SetFilter(-1, 1, timedelta(30), timedelta(45))
            self.option_symbol[ticker] = option.Symbol

            self.atr[ticker] = self.ATR(equity, 14, Resolution.Daily)
            self.macd[ticker] = self.MACD(equity, 12, 26, 9, MovingAverageType.Wilders, Resolution.Daily, Field.Close)
            self.rsi[ticker] = self.RSI(equity, 14, MovingAverageType.Wilders, Resolution.Daily)

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

        for ticker in self.tickers:
            if self.HasOpenStraddle(ticker):
                self._exit_checks(ticker)
                continue

            if self.lastEntryTime.get(ticker) and (self.Time - self.lastEntryTime[ticker]).days < self.cooldown_days:
                continue

            if self._atr_breakout(ticker) and self._macd_cross(ticker) and self._rsi_neutral(ticker):
                self._enter_straddle(ticker, slice)

    def _atr_breakout(self, ticker):
        if not self.atr[ticker].IsReady:
            return False
        price = self.Securities[ticker].Close
        low = self.Securities[ticker].Low
        atr_val = self.atr[ticker].Current.Value
        return (price - low) > self.atrMultiplier * atr_val

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

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

    def _enter_straddle(self, ticker, slice):
        chain = slice.OptionChains.get(self.option_symbol[ticker])
        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[ticker].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[ticker] = totalCost
        self.entryExpiry[ticker] = expiry
        self.entrySpot[ticker] = spot
        self.lastEntryTime[ticker] = self.Time

        self.Debug(f"Entered straddle: {ticker} | Spot: {spot:.2f}, Strike: {atm_strike}, Qty: {qty}")

    def _exit_checks(self, ticker):
        positions = [h for h in self.Portfolio.Values if h.Invested and h.Symbol.SecurityType == SecurityType.Option and h.Symbol.Underlying.Value == ticker]
        if not positions:
            return

        unreal = sum(h.UnrealizedProfit for h in positions)
        cost = self.entryCost.get(ticker, 1)
        pl_pct = unreal / cost
        days_to_expiry = (self.entryExpiry[ticker].date() - self.Time.date()).days

        if pl_pct >= self.takeProfit or pl_pct <= self.stopLoss or days_to_expiry <= 1:
            self.Debug(f"Exiting {ticker} | PnL: {pl_pct:.2%}, DTE: {days_to_expiry}")
            self.Liquidate(ticker)

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