from AlgorithmImports import *
from datetime import time, datetime, timedelta
from QuantConnect.DataSource import *

class MidPriceFillModel(ImmediateFillModel):
    def __init__(self):
        super().__init__()

    def MarketFill(self, asset: Security, order: MarketOrder) -> OrderEvent:
        if isinstance(order, ComboMarketOrder):
            return self.FillComboMarketOrder(asset, order)
        else:
            return super().MarketFill(asset, order)

    def LimitFill(self, asset: Security, order: LimitOrder) -> OrderEvent:
        if isinstance(order, ComboLimitOrder):
            return self.FillComboLimitOrder(asset, order)
        else:
            return super().LimitFill(asset, order)

    def FillComboMarketOrder(self, asset: Security, order: ComboMarketOrder) -> OrderEvent:
        self.Debug("Retrieving Custom Fill for ComboMarket")
        mid_price = (asset.BidPrice + asset.AskPrice) / 2
        return OrderEvent(
            order.Id, asset.Symbol, order.Time, order.Status, mid_price, order.Quantity, 0, "Filled at MidPrice"
        )

    def FillComboLimitOrder(self, asset: Security, order: ComboLimitOrder) -> OrderEvent:
        self.Debug("Retrieving Custom Fill for ComboLimit")
        mid_price = (asset.BidPrice + asset.AskPrice) / 2
        if mid_price <= order.LimitPrice:
            return OrderEvent(
                order.Id, asset.Symbol, order.Time, OrderStatus.Filled, mid_price, order.Quantity, 0, "Filled at MidPrice"
            )
        else:
            return OrderEvent(
                order.Id, asset.Symbol, order.Time, OrderStatus.Submitted, mid_price, order.Quantity, 0, "No Fill"
            )


class SPXWOptionsChainAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 6, 1)
        self.SetEndDate(2024,8,1)
        self.SetCash(20000)

        #Trading parameters
        self.min_expiry = 4
        self.max_expiry = 4
        self.strike_distance = 5
        self.delta = 0.55
        self.profit_percentage = 0.1
        self.legs = []
        self.closelegs = []
        self.filled = 0
        self.fill = 0

        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)

        self.spy = self.AddEquity("SPY")
        self.index_symbol = self.AddIndex('SPX', Resolution.Minute).Symbol
        self.option = self.AddIndexOption(self.index_symbol, "SPXW", Resolution.Minute)

        self.option.SetFilter(lambda x: x.Strikes(-5, 5)
                              .Expiration(self.min_expiry, self.max_expiry)
                              .IncludeWeeklys())

        # Set the custom fill model
        self.option.SetFillModel(MidPriceFillModel())

        # Initialize option legs
        self.long_leg = None
        self.short_leg = None

        self.Schedule.On(self.DateRules.WeekStart("SPY"), self.TimeRules.AfterMarketOpen("SPY", 10), self.tradeoptions)

    def CustomSecurityInitializer(self, security):
        security.SetFeeModel(ZeroFeeModel())

    def tradeoptions(self):
        self.RetrieveOptions(3)

    def RetrieveOptions(self, dte):
        self.Log(f"Retrieving options chain for SPXW on {self.Time}")
        chain = self.CurrentSlice.OptionChains.get(self.option.Symbol.Canonical)
        if chain:
            puts, calls = self.GetSortedOptionsByDelta(chain, dte)
            short_leg, long_leg = self.SelectLegs(puts, calls)

            if short_leg and long_leg:
                self.PlaceSpread(short_leg, long_leg)
            else:
                self.Debug("Not all legs were selected for spread.")

    def GetSortedOptionsByDelta(self, chain, dte):
        puts = sorted([contract for contract in chain if contract.Right == OptionRight.Put and (contract.Expiry - self.Time).days == dte],
                      key=lambda x: (abs(x.Greeks.Delta + self.delta) if x.Greeks else float('inf'), (x.Expiry - self.Time).days))
        calls = sorted([contract for contract in chain if contract.Right == OptionRight.Call and (contract.Expiry - self.Time).days == dte],
                       key=lambda x: (abs(x.Greeks.Delta - self.delta) if x.Greeks else float('inf'), (x.Expiry - self.Time).days))
        return puts, calls

    def SelectLegs(self, puts, calls):
        short_leg = calls[0] if calls else self.Debug("long leg not available")
        long_leg = self.FindOptionByStrike(calls, short_leg.Strike - 5) if short_leg else None
        self.Debug(f"Legs Selected: {long_leg}|{short_leg}")
        return short_leg, long_leg

    def PlaceSpread(self, short_leg, long_leg):
        self.filled = 0
        self.fill = 0
        self.legs = []
        self.closelegs = []
        self.long_leg = long_leg
        self.short_leg = short_leg

        contracts = 1

        self.legs = [
            Leg.Create(short_leg.Symbol, -1),
            Leg.Create(long_leg.Symbol, 1)
        ]
        self.closelegs = [
            Leg.Create(short_leg.Symbol, -1),
            Leg.Create(long_leg.Symbol, 1)
        ]
        self.ComboMarketOrder(self.legs, contracts, tag='OpenSpread')

        self.Debug(f"Spread open order: {contracts} contracts on {self.Time}")

    def FindOptionByStrike(self, options, strike):
        return next((option for option in options if option.Strike == strike), None)

    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        if orderEvent.Status == OrderStatus.Filled and order and 'OpenSpread' in order.Tag:
            self.Debug(f"Order Filled: Tag=OpenSpread, {orderEvent.FillQuantity} @ {orderEvent.FillPrice}, Order ID: {orderEvent.OrderId}")
            if orderEvent.FillQuantity > 0:
                self.filled += 1
                self.fill += orderEvent.FillPrice
            if orderEvent.FillQuantity < 0:
                self.filled += 1
                self.fill -= orderEvent.FillPrice
            if self.filled == 2:
                absfill = abs(self.fill)
                limit = ((5 - absfill) * self.profit_percentage) + absfill
                qty = -abs(orderEvent.FillQuantity)
                self.ComboLimitOrder(self.closelegs, qty, limit, tag='CloseLimitIC')
                self.Debug(f"Filled: {self.fill}. Close Limit {limit}.")

    def OnData(self, slice: Slice):
        pass