from datetime import timedelta, time, datetime
from AlgorithmImports import *

class RapidDropReboundAlgorithm(QCAlgorithm):

    def Initialize(self):
        # --------------------- USER CONFIGURATION: BACKTEST SETUP ---------------------
        self.SetStartDate(2025, 1, 2)  # Backtest start date (YYYY, M, D)
        self.SetEndDate(2025, 5, 31)   # Backtest end date (YYYY, M, D)
        self.SetCash(100000)           # Initial cash for backtesting

        # Time constraints for entries
        self.entry_start_time = time(10, 0)  # can open trades starting at 10:00
        self.entry_end_time   = time(15, 0)  # can open trades up to 15:00

        # --------------------- USER CONFIGURATION: TICKERS ---------------------
        tickers = ["SOXL", "SOXS", "FNGD", "DRV"]
        # --------------------- END OF USER CONFIGURATION: TICKERS ---------------------

        # Retrieve parameters if provided; otherwise use defaults
        self.drop_threshold       = float(self.GetParameter("drop_threshold"))       if self.GetParameter("drop_threshold")       else 0.0275
        self.entry_multiplier     = float(self.GetParameter("entry_multiplier"))     if self.GetParameter("entry_multiplier")     else 0.2
        self.exit_multiplier      = float(self.GetParameter("exit_multiplier"))      if self.GetParameter("exit_multiplier")      else 0.8
        self.stop_loss_multiplier = float(self.GetParameter("stop_loss_multiplier")) if self.GetParameter("stop_loss_multiplier") else 0.01
        drop_window_minutes       = int(self.GetParameter("drop_window_minutes"))    if self.GetParameter("drop_window_minutes")  else 15

        # Optional parameter to disable EOD selling
        self.disable_selling_before_close = True

        # Convert from minutes to timedelta
        self.drop_window = timedelta(minutes=drop_window_minutes)

        #
        # -------------- NEW: MODE VIA INTEGER PARAMETER ---------------
        #
        # We map: 0 -> "AUTOMATIC", 1 -> "CASH", 2 -> "MARGIN_LOW_25K", 3 -> "MARGIN_ABOVE_25K"
        # If user provides something invalid or no value, default to "AUTOMATIC".
        #
        mode_map = {0: "AUTOMATIC", 1: "CASH", 2: "MARGIN_LOW_25K", 3: "MARGIN_ABOVE_25K"}
        mode_int_param = self.GetParameter("mode")  # The integer param

        if not mode_int_param:
            # Default if not specified
            self.mode_param = "AUTOMATIC"
        else:
            try:
                mode_int = int(mode_int_param)
                self.mode_param = mode_map.get(mode_int, "AUTOMATIC")
            except:
                self.mode_param = "AUTOMATIC"
        #
        # -------------- END NEW PARAM HANDLING ---------------
        #

        # Decide final mode
        self.mode = self.DetermineMode()
        self.Debug(f"Selected Mode = {self.mode}")

        # Data structures for the strategy
        self.symbols = {}
        self.symbol_states = {}
        for ticker in tickers:
            symbol = self.AddEquity(ticker, Resolution.Minute).Symbol
            self.symbols[ticker] = symbol
            self.symbol_states[symbol] = {
                "max_price": None,
                "time_of_max": None,
                "drop_event_active": False,
                "initial_price": None,
                "bottom_price": None,
                "has_entered": False,
                "drop_percent": 0
            }

        # Only one trade at a time
        self.current_trade_symbol = None

        # For daily state reset
        self.last_date = None
        self.pending_reset = False

        # Schedule events
        reset_minutes_before_close = int(self.GetParameter("reset_minutes_before_close")) if self.GetParameter("reset_minutes_before_close") else 5
        close_minutes_before_close = int(self.GetParameter("close_minutes_before_close")) if self.GetParameter("close_minutes_before_close") else 5

        # Daily reset event
        self.Schedule.On(
            self.DateRules.EveryDay(),
            self.TimeRules.BeforeMarketClose(tickers[0], -reset_minutes_before_close),
            self.OnSchedule
        )

        # Daily close/liquidate event
        self.Schedule.On(
            self.DateRules.EveryDay(),
            self.TimeRules.BeforeMarketClose(tickers[0], close_minutes_before_close),
            self.close_positions
        )

        # ----------------- NEW DATA STRUCTURES FOR SIMPLE RESTRICTION CHECKS -----------------
        # Track day trades (opened & closed same day) over last 5 trading days
        self.dayTradeDates = []  # each entry is a date of the day-trade
        # Track unsettled funds (extremely simplified: if we SELL, we consider those funds unsettled for 2 days)
        # We'll store (symbol, settle_datetime)
        self.unsettledSales = {}
        # Track position entry time to detect if it closes the same day (for day trade count)
        self.position_open_time = {}


    # --------------------------------- MODE / AUTOMATIC DETECTION ---------------------------------
    def DetermineMode(self):
        """
        Select final mode based on the user-provided integer => string mapping.
        If "AUTOMATIC", attempt to detect from brokerage account type + total portfolio value.
        """
        mode = self.mode_param.upper()

        # If user didn't specify or specifically set "AUTOMATIC"...
        if mode == "AUTOMATIC":
            # Attempt to read from BrokerageModel
            try:
                if self.BrokerageModel.AccountType == AccountType.Cash:
                    return "CASH"
                else:
                    # It's a margin account; check total portfolio value
                    if self.Portfolio.TotalPortfolioValue < 25000:
                        return "MARGIN_LOW_25K"
                    else:
                        return "MARGIN_ABOVE_25K"
            except:
                # If we can't read brokerage model for some reason, fallback
                if self.Portfolio.TotalPortfolioValue < 25000:
                    return "MARGIN_LOW_25K"
                else:
                    return "MARGIN_ABOVE_25K"

        # Otherwise, just return what the user set
        return mode

    # ----------------------------------------------------------------------------------------------
    def OnData(self, data):
        """ Main event loop for handling price data updates and strategy logic. """

        # 1) Check if we've started a new trading day, reset daily states if needed
        if self.last_date != self.Time.date():
            self.last_date = self.Time.date()
            self.ResetDailyState()  # normal daily reset if there's no pending carryover

        # 2) If we previously skipped a reset (carried over trade), check if we can do it now
        if self.pending_reset and not self.Portfolio.Invested:
            self.Debug("Carried-over trade just closed. Performing the skipped daily reset now.")
            self.ResetAllSymbolStates()
            self.pending_reset = False

        # 3) Also, on each new bar, drop from dayTradeDates anything older than 5 trading days
        self.PruneDayTrades()

        # 4) Prune any unsettled sales that have “settled” (older than 2 days)
        self.PruneUnsettledSales()

        # 5) Strategy logic
        current_time = self.Time.time()
        in_entry_window = (current_time >= self.entry_start_time and current_time <= self.entry_end_time)

        for ticker, symbol in self.symbols.items():
            # Ensure we have fresh data for the symbol
            if not data.Bars.ContainsKey(symbol):
                continue

            price = data.Bars[symbol].Close
            state = self.symbol_states[symbol]

            # If no one is invested, track max price for a potential drop
            # Only track if not in an active drop event and we are in the entry window
            if not self.Portfolio.Invested and not state["drop_event_active"] and in_entry_window:
                # Update max price or reset if old
                if state["max_price"] is None or price > state["max_price"]:
                    state["max_price"] = price
                    state["time_of_max"] = self.Time
                else:
                    if self.Time - state["time_of_max"] > self.drop_window:
                        state["max_price"] = price
                        state["time_of_max"] = self.Time

                # Check if we got a rapid drop
                if state["max_price"] is not None:
                    drop_percent = (state["max_price"] - price) / state["max_price"]
                    # Drop must be within drop_window
                    if drop_percent >= self.drop_threshold and (self.Time - state["time_of_max"]) <= self.drop_window:
                        state["drop_event_active"] = True
                        state["initial_price"]     = state["max_price"]
                        state["bottom_price"]      = price
                        state["drop_percent"]      = drop_percent
                        self.Debug(f"{symbol} rapid drop detected. Initial: {state['initial_price']:.2f}, "
                                   f"Price: {price:.2f}, Drop: {drop_percent:.2f}, Threshold: {self.drop_threshold}")

            elif not self.Portfolio.Invested and state["drop_event_active"]:
                # Entry Condition
                if not state["has_entered"] and in_entry_window and (self.current_trade_symbol is None):
                    entry_threshold = state["bottom_price"] * (1 + (self.entry_multiplier * state["drop_percent"]))
                    if price >= entry_threshold:
                        # -------------- CHECK RESTRICTIONS BEFORE WE BUY ---------------
                        if self.CanOpenPosition():
                            self.SetHoldings(symbol, 1.0, tag=f"BUY {symbol} at {price:.2f}")
                            state["has_entered"] = True
                            self.current_trade_symbol = symbol
                            # Record the time we opened to detect day trade if closed same date
                            self.position_open_time[symbol] = self.Time
                            self.Debug(f"BUY {symbol} at {price:.2f} (Rebounded to {entry_threshold:.2f})")
                        else:
                            self.Debug(f"Skipping BUY for {symbol} due to trading restrictions.")

            elif self.Portfolio.Invested and state["drop_event_active"]:
                # If we've entered a position, check exit conditions
                if state["has_entered"] and self.current_trade_symbol == symbol:
                    exit_threshold = state["bottom_price"] * (1 + (self.exit_multiplier * state["drop_percent"]))

                    # Exit Condition
                    if price >= exit_threshold:
                        self.Liquidate(symbol, tag=f"SELL {symbol} at {price:.2f} (Reached exit threshold)")
                        self.Debug(f"SELL {symbol} at {price:.2f} (Reached exit threshold {exit_threshold:.2f})")
                        self.ResetSymbolState(symbol)
                        self.current_trade_symbol = None

                    # Stop Loss Condition
                    entry_price = self.Portfolio[symbol].AveragePrice
                    stop_loss_threshold = entry_price * (1 - self.stop_loss_multiplier)
                    if price <= stop_loss_threshold:
                        self.Liquidate(symbol, tag=f"SL SELL {symbol} at {price:.2f} (Stop loss)")
                        self.Debug(f"SL SELL {symbol} at {price:.2f} (Reached stop loss threshold {stop_loss_threshold:.2f})")
                        self.ResetSymbolState(symbol)
                        self.current_trade_symbol = None

    # -------------------------------------------------------------
    # VERY SIMPLIFIED FREE RIDING / PDT CHECKS
    # -------------------------------------------------------------
    def CanOpenPosition(self):
        """Returns True if we are allowed to open a new position under the current mode."""
        # 1) If CASH mode, disallow opening if we have any unsettled sales.
        if self.mode == "CASH":
            # If we still have unsettled funds, skip
            if len(self.unsettledSales) > 0:
                self.Debug("Cannot open new position in CASH mode because prior sales haven't settled (T+2).")
                return False
            return True

        # 2) If MARGIN < 25K, check day trade count in last 5 days
        if self.mode == "MARGIN_LOW_25K":
            # We already pruned old day trades, so self.dayTradeDates is only the last 5 business days
            # If we already have 4 or more, skip new intraday trades
            self.debug(f"Len day trades {len(self.dayTradeDates)}")
            if len(self.dayTradeDates) >= 4:
                self.Debug("Cannot open new intraday position in MARGIN_LOW_25K mode; 4+ day trades in last 5 days.")
                return False
            return True

        # 3) If MARGIN >= 25K or anything else, no additional checks
        return True


    def on_order_event(self, order_event: OrderEvent) -> None:
        """Capture order fills and track day trades and unsettled funds."""
        if order_event.Status != OrderStatus.Filled:
            return

        # Basic info
        symbol = order_event.Symbol
        order = self.Transactions.GetOrderById(order_event.OrderId)

        # If it's an opening BUY, just note the time
        if order.Direction == OrderDirection.Buy:
            self.position_open_time[symbol] = self.Time
            return

        # If it's a SELL, mark unsettled if we're in CASH mode; check for day trades
        if order.Direction == OrderDirection.Sell:
            # 1) If CASH mode, mark T+2 settlement
            if self.mode == "CASH":
                settle_time = self.Time + timedelta(days=2)
                self.unsettledSales[symbol] = settle_time
                self.Debug(f"Marked {symbol} sale as unsettled until {settle_time} for CASH mode (T+2).")

            # 2) Check if day trade (opened and closed same date)
            open_time = self.position_open_time.get(symbol, None)
            if open_time is not None:
                if open_time.date() == self.Time.date():
                    # This is a day trade
                    self.dayTradeDates.append(self.Time.date())
                    self.Debug(f"Day trade detected for {symbol} on {self.Time.date()}")
                # Remove the open-time so if they later re-open, it won't double count
                del self.position_open_time[symbol]


    def PruneDayTrades(self):
        """
        Remove day-trade entries older than 5 trading days.
        For simplicity, we treat 5 calendar days as a proxy.
        """
        cutoff = self.Time.date() - timedelta(days=5)
        self.dayTradeDates = [d for d in self.dayTradeDates if d >= cutoff]


    def PruneUnsettledSales(self):
        """
        Remove unsettled sales that have passed T+2 settlement date.
        """
        if self.mode == "CASH":
            remove_list = []
            for sym, settle_time in self.unsettledSales.items():
                if self.Time >= settle_time:
                    remove_list.append(sym)
            for sym in remove_list:
                self.Debug(f"Sale for {sym} just settled.")
                del self.unsettledSales[sym]


    # -------------------------------------------------------------
    # END OF STRATEGY, SCHEDULED EVENTS, ETC.
    # -------------------------------------------------------------
    def OnEndOfDay(self):
        # Typically not used for forced liquidation (we use scheduled close_positions)
        pass

    def OnSchedule(self):
        """Fires at (market-close - reset_minutes_before_close)."""
        if not self.Portfolio.Invested or not self.disable_selling_before_close:
            self.Debug("Scheduled reset triggered. Performing daily reset.")
            self.ResetAllSymbolStates()
            self.current_trade_symbol = None
            self.pending_reset = False
        else:
            self.Debug("Scheduled reset triggered, but an open position remains and 'disable_selling_before_close' is True. Skipping reset now.")
            self.pending_reset = True

    def close_positions(self):
        """Fires at (market-close - close_minutes_before_close)."""
        if not self.disable_selling_before_close:
            if self.Portfolio.Invested:
                self.Liquidate()
                self.Debug("Liquidated positions at end of day (default behavior).")
        else:
            self.Debug("'disable_selling_before_close' = True. Skipping end-of-day liquidation.")

    def ResetAllSymbolStates(self):
        """Helper method to reset all symbols' drop-event-related state."""
        for symbol in self.symbol_states:
            self.symbol_states[symbol] = {
                "max_price": None,
                "time_of_max": None,
                "drop_event_active": False,
                "initial_price": None,
                "bottom_price": None,
                "has_entered": False,
                "drop_percent": 0
            }
        self.Debug("Reset all symbol states.")

    def ResetSymbolState(self, symbol):
        """Reset a single symbol's state after it closes a trade."""
        self.symbol_states[symbol] = {
            "max_price": None,
            "time_of_max": None,
            "drop_event_active": False,
            "initial_price": None,
            "bottom_price": None,
            "has_entered": False,
            "drop_percent": 0
        }

    def ResetDailyState(self):
        """Reset daily states for a new day (unless a position is carried over)."""
        if not self.Portfolio.Invested:
            self.ResetAllSymbolStates()
            self.current_trade_symbol = None
            self.pending_reset = False
            self.Debug("New day. No open positions. Daily state reset.")
        else:
            self.Debug("New day. There's a carried-over position. Will skip daily reset here.")