from AlgorithmImports import *
from datetime import timedelta, datetime

class SupportResistanceOptionsTradingAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 1, 1)
        self.SetEndDate(2024, 6, 29)
        self.SetCash(100000)

        self.symbols = ["SPY"]
        self.trade_amount = self.Portfolio.Cash  # Allocate all available cash to trades

        for symbol in self.symbols:
            equity = self.AddEquity(symbol, Resolution.Minute)
            equity.SetLeverage(1)
            option = self.AddOption(symbol, Resolution.Minute)
            option.SetFilter(-10, 10, timedelta(0), timedelta(30))

        self.support_levels = {}
        self.resistance_levels = {}
        self.entry_contracts = {}
        self.traded_today = {symbol: False for symbol in self.symbols}
        self.open_orders = {}

        for symbol in self.symbols:
            self.Schedule.On(self.DateRules.EveryDay(symbol), self.TimeRules.AfterMarketOpen(symbol, 30), 
                             Action(lambda symbol=symbol: self.IdentifySupportResistanceLevels(symbol)))
            self.Schedule.On(self.DateRules.EveryDay(symbol), self.TimeRules.BeforeMarketClose(symbol, 10), 
                             Action(lambda symbol=symbol: self.CloseExpiringPositions(symbol)))
            self.Schedule.On(self.DateRules.EveryDay(symbol), self.TimeRules.BeforeMarketClose(symbol, 10), 
                             Action(self.LiquidateUnexpectedShares))

            # Schedule to reset the traded_today flag after market open
            self.Schedule.On(self.DateRules.EveryDay(symbol), self.TimeRules.AfterMarketOpen(symbol, 0),
                             Action(lambda symbol=symbol: self.ResetTradedTodayFlag(symbol)))

    def ResetTradedTodayFlag(self, symbol):
        self.traded_today[symbol] = False
        self.Debug(f"{symbol}: Reset traded_today flag for the new day.")

    def OnData(self, data):
        if self.IsWarmingUp:
            return

        for symbol in self.symbols:
            if symbol in data.Bars and not self.traded_today[symbol]:
                if self.IsMarketOpen(symbol):
                    if self.IsSupportConfirmed(symbol):
                        self.Debug(f"{symbol}: Support confirmed. Executing trade.")
                        self.ExecuteTrade(symbol)

    def IdentifySupportResistanceLevels(self, symbol):
        history = self.History([symbol], 60, Resolution.Minute)
        if not history.empty:
            data_4_hour = history.loc[symbol].resample('4H').agg({'open': 'first', 'high': 'max', 'low': 'min', 'close': 'last'})
            data_daily = history.loc[symbol].resample('D').agg({'open': 'first', 'high': 'max', 'low': 'min', 'close': 'last'})

            self.support_levels[symbol] = min(data_4_hour['low'].min(), data_daily['low'].min())
            self.resistance_levels[symbol] = max(data_4_hour['high'].max(), data_daily['high'].max())

            self.Debug(f"{symbol} Support Levels: {self.support_levels[symbol]}")
            self.Debug(f"{symbol} Resistance Levels: {self.resistance_levels[symbol]}")

    def IsSupportConfirmed(self, symbol):
        current_price = self.Securities[symbol].Price
        support_level = self.support_levels.get(symbol, None)
        return current_price >= support_level if support_level is not None else False

    def ExecuteTrade(self, symbol):
        # Ensure only one trade per day per symbol
        if self.traded_today[symbol]:
            self.Debug(f"{symbol}: Already traded today.")
            return

        contract = self.GetClosestToResistanceOptionContract(symbol, OptionRight.Call)
        if contract and self.Securities.ContainsKey(contract):
            if not self.Securities[contract].Invested and self.Securities[contract].HasData:
                price = self.Securities[contract].Price
                
                # Ensure price is non-zero and contract is liquid
                if price > 0 and self.Securities[contract].Volume > 0:
                    quantity = int(self.Portfolio.Cash // (price * 100))  # Allocate all cash to trade
                    if quantity > 0:
                        self.MarketOrder(contract, quantity)
                        self.entry_contracts[symbol] = contract
                        self.traded_today[symbol] = True  # Set flag to prevent further trades for today
                        self.open_orders[contract] = quantity  # Track the open order
                        self.Debug(f"Executed trade for {symbol}, contract: {contract}, quantity: {quantity}, price: {price}")
                    else:
                        self.Debug(f"{symbol}: Quantity calculated as zero, trade not executed.")
                else:
                    self.Debug(f"{symbol}: Price is zero or volume is too low, trade not executed.")
            else:
                self.Debug(f"{symbol}: No valid contract or already invested.")
        else:
            self.Debug(f"{symbol}: No suitable contract found.")

    def GetClosestToResistanceOptionContract(self, symbol, option_right):
        contracts = self.OptionChainProvider.GetOptionContractList(symbol, self.Time)
        if not contracts:
            self.Debug(f"No contracts found for {symbol}")
            return None

        resistance_level = self.resistance_levels.get(symbol, None)
        if resistance_level is None or resistance_level == 0:
            self.Debug(f"No resistance level found or it's zero for {symbol}")
            return None

        # Filter for liquid contracts with sufficient volume and non-zero price
        liquid_contracts = [
            contract for contract in contracts 
            if contract.ID.OptionRight == option_right 
            and contract.ID.Date > self.Time  # Ensure it's not expiring today
            and self.Securities.ContainsKey(contract) 
            and self.Securities[contract].Price > 0  # Ensure non-zero price
            and self.Securities[contract].Volume > 0  # Ensure it has volume
        ]

        if not liquid_contracts:
            self.Debug(f"No liquid contracts found for {symbol}")
            return None

        # Find the contract with the strike price closest to the resistance level
        closest_contract = min(
            liquid_contracts,
            key=lambda x: abs(x.ID.StrikePrice - resistance_level),
            default=None
        )

        if closest_contract:
            self.Debug(f"Selected contract closest to resistance for {symbol}: Strike: {closest_contract.ID.StrikePrice}, Expiry: {closest_contract.ID.Date}, Right: {closest_contract.ID.OptionRight}")
            return closest_contract

        self.Debug(f"No suitable contract found for {symbol}")
        return None

    def CloseExpiringPositions(self, symbol):
        if symbol not in self.entry_contracts:
            self.Debug(f"{symbol}: No entry contracts found.")
            return

        contract = self.entry_contracts[symbol]
        if contract.ID.Date <= self.Time + timedelta(days=1):
            if contract in self.Portfolio and self.Portfolio[contract].Invested:
                quantity = self.Portfolio[contract].Quantity
                self.MarketOrder(contract, -quantity)
                self.traded_today[symbol] = False
                if contract in self.open_orders:
                    del self.open_orders[contract]
                self.Debug(f"Closed expiring position for {symbol}, contract: {contract}, quantity: {quantity}")

    def LiquidateUnexpectedShares(self):
        for security in self.Portfolio.Values:
            if security.Invested and security.Type == SecurityType.Equity:
                self.Debug(f"Unexpected shares found for {security.Symbol}. Liquidating...")
                self.Liquidate(security.Symbol)

    def OnOrderEvent(self, orderEvent):
        self.Debug(str(orderEvent))