from AlgorithmImports import *


class VolatilityStraddleAlgorithm(QCAlgorithm):
    """ATM long straddle template for real option backtests in QuantConnect/LEAN."""

    def Initialize(self):
        self.SetStartDate(2022, 1, 1)
        self.SetEndDate(2024, 1, 1)
        self.SetCash(100000)

        self.ticker = "SPY"
        self.target_dte = 30
        self.holding_days = 5
        self.entry_every_days = 5
        self.contract_quantity = 1

        equity = self.AddEquity(self.ticker, Resolution.Minute)
        option = self.AddOption(self.ticker, Resolution.Minute)
        option.SetFilter(self.OptionFilter)

        self.underlying = equity.Symbol
        self.option_symbol = option.Symbol
        self.next_entry_time = self.StartDate
        self.open_groups = []

    def OptionFilter(self, universe):
        min_dte = max(1, self.target_dte - 10)
        max_dte = self.target_dte + 10
        return universe.IncludeWeeklys().Strikes(-10, 10).Expiration(min_dte, max_dte)

    def OnData(self, slice):
        self.CloseExpiredHoldingGroups()

        if self.Time < self.next_entry_time:
            return

        chain = slice.OptionChains.get(self.option_symbol)
        if chain is None:
            return

        contracts = [contract for contract in chain if contract.Expiry.date() > self.Time.date()]
        if not contracts:
            return

        expiry = min(contracts, key=lambda contract: abs((contract.Expiry.date() - self.Time.date()).days - self.target_dte)).Expiry
        expiry_contracts = [contract for contract in contracts if contract.Expiry == expiry]
        spot = self.Securities[self.underlying].Price

        calls = [contract for contract in expiry_contracts if contract.Right == OptionRight.Call]
        puts = [contract for contract in expiry_contracts if contract.Right == OptionRight.Put]
        if not calls or not puts:
            return

        call = min(calls, key=lambda contract: abs(contract.Strike - spot))
        put = min(puts, key=lambda contract: abs(contract.Strike - spot))

        self.MarketOrder(call.Symbol, self.contract_quantity)
        self.MarketOrder(put.Symbol, self.contract_quantity)

        self.open_groups.append(
            {
                "entry_time": self.Time,
                "exit_time": self.Time + timedelta(days=self.holding_days),
                "symbols": [call.Symbol, put.Symbol],
            }
        )
        self.next_entry_time = self.Time + timedelta(days=self.entry_every_days)

        self.Debug(
            f"Opened ATM straddle {call.Symbol.Value}, {put.Symbol.Value}; "
            f"spot={spot:.2f}; expiry={expiry.date()}"
        )

    def CloseExpiredHoldingGroups(self):
        remaining_groups = []
        for group in self.open_groups:
            if self.Time < group["exit_time"]:
                remaining_groups.append(group)
                continue

            for symbol in group["symbols"]:
                holding = self.Portfolio[symbol]
                if holding.Invested:
                    self.MarketOrder(symbol, -holding.Quantity)
            self.Debug(f"Closed straddle group from {group['entry_time']}")

        self.open_groups = remaining_groups

    def OnEndOfAlgorithm(self):
        self.Debug(f"Final portfolio value: {self.Portfolio.TotalPortfolioValue:.2f}")