# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#region imports
from AlgorithmImports import *
#endregion
import pandas_market_calendars as mcal
from datetime import datetime, timedelta


class IndexOptionIronCondorAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2025, 1, 1)
        self.set_end_date(2025, 2, 28)
        self.set_cash(1000000)

        index = self.add_index("SPX", Resolution.MINUTE).symbol
        option = self.add_index_option(index, "SPXW", Resolution.MINUTE)
        option.set_filter(lambda x: x.weeklys_only().strikes(-50, 50).expiration(5, 18))
        self.spxw = option.symbol

        self.bb = self.bb(index, 10, 2, resolution=Resolution.DAILY)
        self.warm_up_indicator(index, self.bb)

        self.log(f"time_zone: {self.time_zone}")
        self.time_logged = None
        self.last_date = None
        self.last_answer = False

    def is_last_trading_day_of_week(self, date: Union[datetime, str]) -> bool:
        date = datetime.strptime(date, '%Y-%m-%d') if isinstance(date, str) else date

        if self.last_date and self.last_date == date.date():
            return self.last_answer

        self.last_date = date.date()

        nyse = mcal.get_calendar('NYSE')
        # Define the start and end of the week
        start_of_week = date - timedelta(days=date.weekday())
        end_of_week = start_of_week + timedelta(days=4)

        # Get the NYSE schedule for the week
        week_schedule = nyse.schedule(start_date=start_of_week.strftime('%Y-%m-%d'),
                                    end_date=end_of_week.strftime('%Y-%m-%d'))
        week_dates = week_schedule.index.strftime('%Y-%m-%d')

        last_trading_day = week_dates[-1] if week_schedule.shape[0] > 0 else None

        self.last_answer = date.strftime('%Y-%m-%d') == last_trading_day
        return self.last_answer

    def next_trading_day_of_week(self, date: Union[datetime, str]) -> bool:
        date = datetime.strptime(date, '%Y-%m-%d') if isinstance(date, str) else date

        nyse = mcal.get_calendar('NYSE')
        # Define the start and end of the week
        start_of_next_week = date + timedelta(days=7) - timedelta(days=date.weekday())
        end_of_next_week = start_of_next_week + timedelta(days=4)

        # Get the NYSE schedule for the week
        week_schedule = nyse.schedule(start_date=start_of_next_week.strftime('%Y-%m-%d'),
                                    end_date=end_of_next_week.strftime('%Y-%m-%d'))
        week_dates = week_schedule.index.strftime('%Y-%m-%d')

        last_trading_day = week_dates[-1] if week_schedule.shape[0] > 0 else None

        return last_trading_day


    def on_data(self, slice: Slice) -> None:
        

        #if self.portfolio.invested: return

        if not self.is_last_trading_day_of_week(self.Time):
            #if not self.time_logged or self.time_logged and self.time_logged != self.Time.date():
            #    self.time_logged = self.Time.date()
            #    #self.log(f"on_data: {self.time_logged} is not a day to trade")
            return

        if not (self.Time.hour == 15 and self.Time.minute == 50):
            return
        
        #self.log(f"on_data: {self.Time}")

        # Get the OptionChain
        chain = slice.option_chains.get(self.spxw)
        if not chain: return

        # Get the closest expiry date
        expiry = self.next_trading_day_of_week(self.Time) #min([x.expiry for x in chain])
        #for x in chain:
        #    if x.right == OptionRight.PUT:
        #        self.log(f"{x}")
        chain = [x for x in chain if x.expiry.strftime('%Y-%m-%d') == expiry]

        # Separate the call and put contracts and sort by Strike to find OTM contracts
        #calls = sorted([x for x in chain if x.right == OptionRight.CALL], key=lambda x: x.strike, reverse=True)
        puts = sorted([x for x in chain if x.right == OptionRight.PUT], key=lambda x: x.strike)
        #if len(calls) < 3 or len(puts) < 3: return
        if len(puts) < 3: return

        self.log(f"on_data:  processing option {puts[3]}, Underlying: {puts[3].UnderlyingLastPrice} Next expiration date {self.next_trading_day_of_week(self.Time)}")
        self.MarketOrder(puts[3].symbol, -1)
        
        # Create combo order legs
        '''
        price = self.bb.price.current.value
        quantity = 1
        if price > self.bb.upper_band.current.value or price < self.bb.lower_band.current.value:
            quantity = -1

        legs = [
            Leg.create(calls[0].symbol, quantity),
            Leg.create(puts[0].symbol, quantity)
        ]

        self.log(f"on_data: order {calls[0]}")
        self.log(f"               {puts[0]}")
        self.log(f"               {calls[2]}")
        self.log(f"               {puts[2]}")
        #self.log(f"               {*legs[1]}")
        #self.log(f"               {*legs[2]}")
        #self.log(f"               {*legs[3]}")

        self.combo_market_order(legs, 10, asynchronous=True)
        '''