# region imports
from AlgorithmImports import *
from datetime import timedelta
import pandas as pd
# endregion

class SimpleStraddleBacktest(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2020, 1, 1)
        self.set_end_date(2024, 6, 1)
        self.set_cash(100_000)

        # Setting up asset and equity
        self._option = self.add_option("QQQ", Resolution.MINUTE)        
        self._symbol = self._option.symbol
        self._option.set_filter(timedelta(7), timedelta(7))

        # Schedule liquidate
        self.schedule.on(self.date_rules.every_day(self._symbol), self.time_rules.at(9,33), self.liquidate_positions)
        
    def on_data(self, data:Slice):

        # Start the strategy at 3:58 PM
        if self.time.hour == 15 and self.time.minute == 38:
            
            # Find the list of contracts
            chain = data.option_chains.get(self._symbol)

            if chain is None:
                return
            
            calls = [x for x in chain if x.right == OptionRight.CALL]
            puts = [x for x in chain if x.right == OptionRight.PUT]
            
            straddles = [[call, put] for call in calls for put in puts if call.strike == put.strike]
            self.cheapest_straddle = sorted(straddles, key=lambda straddle: straddle[0].last_price + straddle[1].last_price)[0]
            self.log(f"Cheapest straddle contracts: {self.cheapest_straddle[0].symbol.value, self.cheapest_straddle[1].symbol.value}")

            # Set leverage
            self.securities[self.cheapest_straddle[0].symbol].set_leverage(1)
            self.securities[self.cheapest_straddle[1].symbol].set_leverage(1)


            if not self.portfolio.invested:
                # Buying straddle
                self.set_holdings(self.cheapest_straddle[0].symbol, .5)
                self.set_holdings(self.cheapest_straddle[1].symbol, .5)

    def liquidate_positions(self):
        self.liquidate()