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)
        index_symbol = self.AddIndex("SPX").Symbol
        option = self.AddIndexOption(index_symbol, "SPXW")
        self.symbol = option.Symbol
        option.SetFilter(-2, 2, timedelta(0), timedelta(2))
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        option.EnableGreekApproximation = True
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), self.TimeRules.Every(timedelta(seconds=30)), self.Trade)

    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):
            if i < 2:  # Print the first two rows
                self.Debug(str(order))
            trigger_time = datetime.fromtimestamp(int(order['Trigger Time']))
            if self.Time.date() == trigger_time.date():
                expiry = datetime.strptime(str(order['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: {order}")
                    continue
                contracts.sort(key=lambda x: abs(x.ID.StrikePrice - float(order['Strike 1'])))
                quantity = int(order['Order Quantity'])
                if self.Securities.ContainsKey(contracts[0]):
                    contract = self.Securities[contracts[0]]
                    legs = []
                    for i in range(1, 5):
                        if order[f'Right {i}']:
                            right = OptionRight.Call if order[f'Right {i}'] == 'C' else OptionRight.Put
                            legs.append(Leg.Create(contracts[i-1], quantity, (contract.AskPrice + contract.BidPrice) / 2))
                            self.Debug(f"Added leg: {contracts[i-1].Symbol.Value}, {right}, {quantity}")
                    ticket = self.ComboLegLimitOrder(legs, 1)
                    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(str(orderEvent))