from AlgorithmImports import *
from datetime import timedelta, datetime

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

        self.symbols = ["TSLA", "NVDA", "GOOGL", "NFLX", "AMD", "AAPL", "META", "MSFT", "AMZN"]
        for symbol in self.symbols:
            equity = self.AddEquity(symbol, Resolution.Minute, extendedMarketHours=True)
            equity.SetLeverage(1)

            option = self.AddOption(symbol)
            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 = {}
        self.stop_losses = {}
        self.scale_out_thresholds = {}

        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.ScaleOutPositions(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(), self.TimeRules.At(15, 50), self.LiquidateUnexpectedShares)
        self.Schedule.On(self.DateRules.WeekEnd(), self.TimeRules.At(15, 50), self.CloseAllPositions)

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

        for symbol in self.symbols:
            if symbol in data.Bars and not self.traded_today[symbol] and self.Time.weekday() < 4:  # Only trade from Monday to Thursday
                if self.IsMarketOpen(symbol):
                    if self.IsSupportConfirmed(symbol):
                        self.Debug(f"{symbol}: Support confirmed. Executing call trade.")
                        self.ExecuteTrade(symbol, OptionRight.Call)
                    elif self.IsSupportBroken(symbol):
                        self.Debug(f"{symbol}: Support broken. Closing position.")
                        self.ClosePosition(symbol)
                        self.UpdateLevels(symbol, broken_support=True)

            self.CheckStopLoss(symbol)

    def IdentifySupportResistanceLevels(self, symbol):
        history = self.History([symbol], 60, Resolution.Minute, extendedMarketHours=True)
        if not history.empty:
            data_30_min = history.loc[symbol].resample('30T').agg({'open': 'first', 'high': 'max', 'low': 'min', 'close': 'last'})
            self.support_levels[symbol] = data_30_min['low'].min()
            self.resistance_levels[symbol] = data_30_min['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 IsSupportBroken(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 UpdateLevels(self, symbol, broken_support):
        if broken_support:
            self.resistance_levels[symbol] = self.support_levels[symbol]
            self.support_levels[symbol] = None
            self.Debug(f"{symbol}: Support broken. New Resistance Level: {self.resistance_levels[symbol]}")
        else:
            self.support_levels[symbol] = self.resistance_levels[symbol]
            self.resistance_levels[symbol] = None
            self.Debug(f"{symbol}: Resistance broken. New Support Level: {self.support_levels[symbol]}")

    def ExecuteTrade(self, symbol, option_right):
        if self.traded_today[symbol]:
            self.Debug(f"{symbol}: Already traded today.")
            return

        contract = self.GetClosestToResistanceOptionContract(symbol, option_right)
        if contract and self.Securities.ContainsKey(contract):
            if not self.Securities[contract].Invested and self.Securities[contract].HasData:
                position_size = self.CalculatePositionSize()
                price = self.Securities[contract].Price
                if price != 0:
                    quantity = position_size // (price * 100)  # 100 shares per contract
                    if quantity > 0:
                        self.MarketOrder(contract, quantity)
                        self.entry_contracts[symbol] = contract
                        self.traded_today[symbol] = True
                        self.open_orders[contract] = quantity  # Track the open order
                        self.stop_losses[contract] = min(price * 0.5, price - (price - self.support_levels[symbol]))  # 50% stop loss or support break
                        self.scale_out_thresholds[contract] = self.CalculateScaleOutThresholds(contract, quantity)  # Scale out thresholds
                        self.Debug(f"Executed trade for {symbol}, contract: {contract}, quantity: {quantity}")
                    else:
                        self.Debug(f"Quantity calculated as zero for {symbol}, skipping trade.")
                else:
                    self.Debug(f"Price is zero for {symbol}, skipping trade.")

    def CalculateScaleOutThresholds(self, contract, quantity):
        if contract.Underlying not in self.resistance_levels:
            self.Debug(f"{contract.Underlying} not in resistance levels.")
            return []

        price = self.Securities[contract].Price
        increment = (self.resistance_levels[contract.Underlying] - price) / 3
        return [(price + increment * (i + 1), quantity // 3) for i in range(3)]

    def CalculatePositionSize(self):
        available_cash = self.Portfolio.Cash
        position_size = available_cash / len(self.symbols) if len(self.symbols) > 0 else 0
        self.Debug(f"Calculated position size: {position_size}")
        return position_size

    def ClosePosition(self, symbol):
        contract = self.entry_contracts.get(symbol, None)
        if contract and 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]
            if contract in self.stop_losses:
                del self.stop_losses[contract]
            if contract in self.scale_out_thresholds:
                del self.scale_out_thresholds[contract]
            self.Debug(f"Closed position for {symbol}, contract: {contract}, quantity: {quantity}")

    def CheckStopLoss(self, symbol):
        contract = self.entry_contracts.get(symbol, None)
        if contract and contract in self.Portfolio and self.Portfolio[contract].Invested:
            current_price = self.Securities[contract].Price
            stop_loss_price = self.stop_losses.get(contract, None)

            if stop_loss_price is not None and (current_price <= stop_loss_price or self.IsSupportBroken(symbol)):
                self.Debug(f"{symbol}: Stop loss triggered. Closing position.")
                self.ClosePosition(symbol)

    def ScaleOutPositions(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 not in self.scale_out_thresholds:
            self.Debug(f"{symbol}: No scale out thresholds found for {contract}.")
            return

        thresholds = self.scale_out_thresholds[contract]
        for price, qty in thresholds:
            if self.Securities[contract].Price >= price:
                self.Debug(f"Scaling out position for {symbol}, contract: {contract}, at price: {price}")
                self.MarketOrder(contract, -qty)
                thresholds.remove((price, qty))
                break

        if not thresholds:
            self.traded_today[symbol] = False
            self.Debug(f"All scale out thresholds met for {symbol}. Resetting trade status.")

    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:
            self.Debug(f"No resistance level found for {symbol}")
            return None

        closest_contract = min(
            (contract for contract in contracts if contract.ID.OptionRight == option_right and contract.ID.Date.weekday() == 4),
            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]
                if contract in self.stop_losses:
                    del self.stop_losses[contract]
                if contract in self.scale_out_thresholds:
                    del self.scale_out_thresholds[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 CloseAllPositions(self):
        for symbol in self.symbols:
            self.ClosePosition(symbol)

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