from AlgorithmImports import *

class EnergeticRedAnt(QCAlgorithm):
    def Initialize(self):
        # Setup
        self.SetStartDate(2024, 1, 1)
        self.SetEndDate(2024, 1, 15)
        self.SetCash(100_000)
        
        # Data
        self.symbol = self.AddIndex("SPX", Resolution.Hour).Symbol
        option = self.AddIndexOption(self.symbol, "SPXW", resolution=Resolution.Hour)
        self.symbol_option = option.Symbol
        option.SetFilter(minStrike=-10, maxStrike=10, minExpiry=timedelta(days=0), maxExpiry=timedelta(days=1))

        # Init
        self.expected_moves = RollingWindow[float](5)
        self.actual_moves   = RollingWindow[float](5)
        self.long_straddle  = None
        self.short_straddle = None

    def OnData(self, slice: Slice) -> None:
        # Check data
        if not slice.ContainsKey(self.symbol): return

        if slice.OptionChains:
            self.Debug(f"time 2 {self.Time}")
            for chain in slice.OptionChains.Values:
                atmCall, atmPut = self.FindATMOption(chain)
                if atmCall is not None and atmPut is not None:
                    # Get straddle price
                    straddle_open_price = slice[atmPut.Symbol].Open + slice[atmCall.Symbol].Open

                    # Calculate straddle-implied expected move and store
                    expected_move = 0.85 * straddle_open_price
                    self.expected_moves.Add(expected_move)
                    self.Log(f"expected_move: {expected_move}")
                else:
                    return

            # Calculate actual move from SPX prices and store
            actual_move = abs(slice[self.symbol].Close - slice[self.symbol].Open)
            self.actual_moves.Add(actual_move)
            self.Log(f"actual_move {actual_move}")

            # Calculate recent average expected and actual moves
            self.Log(f"expected_moves.IsReady {self.expected_moves.IsReady} actual_moves {self.actual_moves.IsReady}")
            if self.expected_moves.IsReady and self.actual_moves.IsReady:
                ave_expected_move = sum(list(self.expected_moves)) / len(list(self.expected_moves)) # self.expected_moves.Count maybe faster
                ave_actual_move   = sum(list(self.actual_moves)) / len(list(self.actual_moves))
                self.Log(f"Average expected move, average actual move: ave_expected_move {ave_expected_move} ave_actual_move {ave_actual_move}")

                # Check if we already have an invested position
                invested = any([self.Portfolio[option.Symbol].Invested for option in [atmCall, atmPut]])
                self.Debug(f"Invested {invested}")

                # Trading logic
                if not invested:
                    self.Debug(f"ave_actual_move {ave_actual_move} ave_expected_move {ave_expected_move}")
                    if ave_actual_move > ave_expected_move:
                        short_straddle = OptionStrategies.ShortStraddle(self.symbol_option, atmCall.Strike, atmCall.Expiry)
                        self.Buy(short_straddle, 1)
                    else:
                        long_straddle = OptionStrategies.Straddle(self.symbol_option, atmPut.Strike, atmPut.Expiry)
                        self.Buy(long_straddle, 1)
        return

    def FindATMOption(self, optionChain):
        underlyingPrice = self.Securities[self.symbol].Price
        minCallDistance = float('inf')
        minPutDistance = float('inf')
        atmCallOption = None
        atmPutOption = None

        # Iterate through the option chain to find the closest call and put options
        for optionContract in optionChain:
            distance = abs(optionContract.Strike - underlyingPrice)
            
            # Update the closest call option
            if optionContract.Right == OptionRight.Call and distance < minCallDistance:
                atmCallOption = optionContract
                minCallDistance = distance
            
            # Update the closest put option
            elif optionContract.Right == OptionRight.Put and distance < minPutDistance:
                atmPutOption = optionContract
                minPutDistance = distance

        return atmCallOption, atmPutOption