from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.Custom import *
from QuantConnect.Orders import *
from QuantConnect.Securities.Option import OptionPriceModels
from datetime import timedelta, datetime
import csv
import io

class SPXWTradingAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 10, 30)
        self.SetCash(100000)
        spx = self.AddIndex("SPX").Symbol
        option = self.AddIndexOption(spx, "SPXW") # SPXW is the target non-standard contract
        option.SetFilter(lambda universe: universe.IncludeWeeklys().Strikes(-2, 2).Expiration(timedelta(0), timedelta(2))) # Note: IncludeWeeklys and limit to 2 days
        self.symbol = option.Symbol
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        option.EnableGreekApproximation = True
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), self.TimeRules.Every(timedelta(seconds=30)), self.Trade)
        self.last_row = None

        # Log the last row during initialization
        self.Trade()
        if self.last_row is not None:
            self.Debug("Last row during initialization: " + str(self.last_row))

    def Trade(self):
        csv_string = self.Download('https://docs.google.com/spreadsheets/d/1wwadCU8msu6FEUJt1ANoZS2qMO2MWiheARrdm7zaQlM/export?format=csv')
        csv_reader = csv.DictReader(io.StringIO(csv_string))
        for i, order in enumerate(csv_reader):
            self.last_row = order

        if self.last_row is None:
            self.Debug("Last row is None")
            return

        expiry = datetime.strptime(str(self.last_row['TWS Contract Date']), '%Y%m%d')
        optionchain = self.OptionChainProvider.GetOptionContractList(self.symbol, self.Time.date())
        contracts = [i for i in optionchain if i.ID.Date.date() == expiry.date()]
        if len(contracts) < 1:
            self.Debug(f"Not enough option contracts for order: {self.last_row}")
            return
        contracts.sort(key=lambda x: abs(x.ID.StrikePrice - float(self.last_row['Strike 1'])))
        quantity = int(self.last_row['Order Quantity'])
        if self.Securities.ContainsKey(contracts[0]):
            contract = self.Securities[contracts[0]]
            legs = []
            total_limit_price = 0
            for i in range(1, 3):
                if self.last_row[f'Right {i}']:
                    right = OptionRight.Call if self.last_row[f'Right {i}'] == 'C' else OptionRight.Put
                    limit_price = (contract.AskPrice + contract.BidPrice) / 2
                    legs.append(Leg.Create(contracts[i-1], quantity, limit_price))
                    total_limit_price += limit_price
                    self.Debug(f"Added leg: {contracts[i-1].Symbol.Value}, {right}, {quantity}")
            ticket = self.ComboLegLimitOrder(legs, total_limit_price)
            self.Debug(f"Symbol: {ticket.Symbol}; Quantity filled: {ticket.QuantityFilled}; Fill price: {ticket.AverageFillPrice}")

    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            self.Debug("Order filled: " + str(orderEvent))