from AlgorithmImports import *

index_names = ["SPX", "VIX", "NDX"]
option_names = ["SPXW", "VIXW", "NQX"]
index = 0
pos_sizes = [0.0006, 0.005, 0.01, 0.025, 0.05, 0.1, 0.2]  # 0-6
stop_losses = [0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 0]  # 0-6
pos_size = 0
stop_loss = -1
option_type = 0  # 0 - both, 1 - put, 2 - call
trade_start = (11, 5, 0)
trade_end = (0, 0, 0)
start_date = (2022, 1, 1)
end_date = (2023, 2, 14)
cash = 100_000
take_profit = 0


class SPXW(QCAlgorithm):
    def Initialize(self):
        self.Log([f"{x}" for x in self.GetParameters()])
        global pos_size, stop_loss, index, option_type
        pos_size = pos_sizes[int(self.GetParameter("pos_size") or pos_size)]
        stop_loss = stop_losses[int(self.GetParameter("stop_loss") or stop_loss)]
        index = int(self.GetParameter("index") or index)
        option_type = int(self.GetParameter("option_type") or option_type)
        self.SetBrokerageModel(InteractiveBrokersBrokerageModel(AccountType.Margin))

        self.SetStartDate(*start_date)
        self.SetEndDate(*end_date)
        self.SetCash(cash)

        self.index_security = self.AddIndex(index_names[index], Resolution.Minute)
        self.index_security = self.AddEquity("SPY", Resolution.Minute)

        # weekly option SPX contracts
        # index_options = self.AddIndexOption(self.index_security.Symbol, option_names[index])
        index_options = self.AddOption(self.index_security.Symbol, Resolution.Minute)
        # set our strike/expiry filter for this option chain
        index_options.SetFilter(lambda u: (u.Strikes(-25, 25)
                                           .Expiration(0, 0)
                                           .IncludeWeeklys()))

        self.option_symbol = index_options.Symbol
        self.ResetTrade()

        self.Schedule.On(self.DateRules.EveryDay(self.option_symbol), self.TimeRules.At(1, 0, 0), self.ResetTrade)
        self.Schedule.On(self.DateRules.EveryDay(self.option_symbol), self.TimeRules.At(*trade_start), self.StartTrade)
        self.Schedule.On(self.DateRules.EveryDay(self.option_symbol), self.TimeRules.At(*trade_end), self.StopTrade)
        # self.SetWarmUp(90, Resolution.Minute)

        if take_profit:
            self.AddRiskManagement(MaximumUnrealizedProfitPercentPerSecurity(take_profit))

    def ResetTrade(self):
        self.trade = False
        self.day_low = 100000
        self.day_high = 0

    def StartTrade(self):
        self.trade = True

    def StopTrade(self):
        self.trade = False

    def OnData(self, slice):
        current_price = self.index_security.Price
        if current_price < self.day_low:
            self.day_low = current_price
        if current_price > self.day_high:
            self.day_high = current_price

        delta = 0.01

        if not self.trade: return

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

        call = [x for x in chain if x.Right == OptionRight.Call]
        put = [x for x in chain if x.Right == OptionRight.Put]

        # we sort the contracts to find contract with the right delta
        put_contract = sorted(put, key=lambda x: abs(abs(x.Greeks.Delta) - delta))
        call_contract = sorted(call, key=lambda x: abs(x.Greeks.Delta - delta))

        # if found, trade it
        if option_type != 2:
            if len(put_contract) == 0:
                return
            else:
                # if current_price <= 1.001 * self.day_low:
                if current_price <= self.day_low:
                    put_ = put_contract[0].Symbol
                    self.SetHoldings(put_, -pos_size)
                    # self.MarketOrder(put_, -1000 * pos_size)
                    self.trade = False
                    return

        if option_type != 1:
            if len(call_contract) == 0:
                return
            else:
                # if current_price >= 0.999 * self.day_high:
                if current_price >= self.day_high:
                    call_ = call_contract[0].Symbol
                    # self.MarketOrder(call_, -1000 * pos_size)
                    self.SetHoldings(call_, -pos_size)
                    self.trade = False