# Install required packages if not already present
!pip install pandas numpy --quiet

import pandas as pd
import numpy as np
import warnings

# Suppress non-critical runtime warnings (e.g., divide-by-zero in empty slices)
warnings.filterwarnings('ignore')

print("Libraries loaded successfully.")

def get_interval_min_data(np_1m: np.ndarray, np_model_predictions: np.ndarray, index: int, index_1m_datetime: int, index_1m_high: int, index_1m_low: int, index_pred_datetime: int):
    """
    Extract 1-minute candle rows for the prediction interval at `index`.

    The interval spans from `predictions[index].datetime` (inclusive)
    to `predictions[index+1].datetime` (exclusive).

    Parameters
    ----------
    np_1m : np.ndarray
        Full 2D array of 1-minute OHLCV rows.
    np_model_predictions : np.ndarray
        Array of prediction data.
    index : int
        Row index in the predictions array.
    index_1m_datetime : int
        Column index for the datetime field in `np_1m`.
    index_1m_high : int
        Column index for the high price in `np_1m`.
    index_1m_low : int
        Column index for the low price in `np_1m`.
    index_pred_datetime : int
        Column index for the datetime field in `np_model_predictions`.

    Returns
    -------
    tuple : (np_temp, np_temp_high, np_temp_low)
        np_temp      — Full 2D slice of 1-minute rows.
        np_temp_high — High prices for the interval.
        np_temp_low  — Low prices for the interval.
    """
    start_time = np.datetime64(np_model_predictions[index][index_pred_datetime])
    end_time   = np.datetime64(np_model_predictions[index + 1][index_pred_datetime])

    # Use pandas to robustly convert the datetime column from the dataset to np.datetime64[ns]
    market_datetimes = pd.to_datetime(np_1m[:, index_1m_datetime])

    # Boolean mask to filter 1-minute rows within the interval
    mask = (
        (market_datetimes >= start_time) &
        (market_datetimes <  end_time)
    )
    np_1m_indices = np.where(mask)[0]

    np_temp      = np_1m[np_1m_indices]
    np_temp_high = np_1m[np_1m_indices, index_1m_high]
    np_temp_low  = np_1m[np_1m_indices, index_1m_low]

    return np_temp, np_temp_high, np_temp_low

def record_trade(
    array_to_save: list,
    datetime,
    predicted_direction: int,
    action: str,
    buy_price: float,
    sell_price: float,
    current_balance: float,
    pnl: float
) -> None:
    """
    Append a trade record to the internal ledger.

    Parameters
    ----------
    array_to_save : list
        The list to which trade records are appended.
    datetime : np.datetime64
        Timestamp of the trade event.
    predicted_direction : int
        The predicted direction (+1 for long, -1 for short).
    action : str
        Trade action label (e.g., 'buy', 'sell - take_profit').
    buy_price : float
        Entry price.
    sell_price : float
        Exit price (0 if still open).
    current_balance : float
        Account balance after the event.
    pnl : float
        Realized PnL in percentage terms for this event.
    """
    array_to_save.append([
        datetime,
        "long" if predicted_direction > 0 else "short",
        action,
        buy_price,
        sell_price,
        current_balance,
        pnl,
    ])

def find_tp_sl_index(
    current_pred_direction: int,
    take_profit_amount: float,
    stop_loss_amount: float,
    np_temp_high: np.ndarray,
    np_temp_low: np.ndarray
):
    """
    Identify the first minute within the interval where TP or SL is breached.

    For LONG positions:
      - TP is triggered when High >= take_profit_amount
      - SL is triggered when Low  <= stop_loss_amount

    For SHORT positions:
      - TP is triggered when Low  <= take_profit_amount
      - SL is triggered when High >= stop_loss_amount

    Returns
    -------
    tuple : (triggered: bool, index: int, sell_price: float)
        triggered — whether TP or SL was hit.
        index     — minute-level index at which the event occurred.
        sell_price— The price at which the TP/SL was hit.
    """
    if current_pred_direction > 0:  # LONG
        idx_high = np.where(np_temp_high >= take_profit_amount)[0]  # TP hit
        idx_low  = np.where(np_temp_low  <= stop_loss_amount)[0]    # SL hit
    else:                                                            # SHORT
        idx_high = np.where(np_temp_high >= stop_loss_amount)[0]    # SL hit (short)
        idx_low  = np.where(np_temp_low  <= take_profit_amount)[0]  # TP hit (short)

    no_high = len(idx_high) == 0
    no_low  = len(idx_low)  == 0

    sell_price = 0.0
    if no_high and no_low:
        # Neither TP nor SL breached within the interval
        return False, -1, sell_price

    elif not no_high and no_low:
        # Only High threshold breached
        df_index = idx_high[0]
        sell_price = np_temp_high[df_index]
        return True, df_index, sell_price

    elif no_high and not no_low:
        # Only Low threshold breached
        df_index = idx_low[0]
        sell_price = np_temp_low[df_index]
        return True, df_index, sell_price

    else:
        # Both thresholds breached — take the one that occurred first
        if idx_high[0] < idx_low[0]:
            df_index = idx_high[0]
            sell_price = np_temp_high[df_index]
        else:
            df_index = idx_low[0]
            sell_price = np_temp_low[df_index]
        return True, df_index, sell_price

def calculate_pnl(
    buy_price: float,
    sell_price: float,
    previous_pred_direction: int,
    leverage: float,
    transaction_fee_percent: float,
    slippage: float
) -> float:
    """
    Calculate the PnL for a trade.

    Parameters
    ----------
    buy_price : float
        The price at which the position was entered.
    sell_price : float
        The price at which the position was exited.
    previous_pred_direction : int
        The direction of the trade (1 for long, -1 for short).
    leverage : float
        Leverage multiplier.
    transaction_fee_percent : float
        Transaction fee percentage.
    slippage : float
        Slippage percentage.

    Returns
    --------
    float
        The calculated PnL in percentage terms.
    """
    if previous_pred_direction > 0:  # LONG: profit when price rises
        pnl = ((sell_price - buy_price) / buy_price) * 100
    else:                               # SHORT: profit when price falls
        pnl = ((buy_price - sell_price) / buy_price) * 100

    return pnl * leverage - transaction_fee_percent - slippage

def generate_mock_1m_data(n_days: int = 10, base_price: float = 30000.0) -> tuple:
    """
    Generate synthetic 1-minute OHLCV candle data and associated column indices.

    Parameters
    ----------
    n_days : int
        Number of trading days to simulate.
    base_price : float
        Starting price for the random walk simulation.

    Returns
    -------
    tuple : (np.ndarray, int, int, int, int)
        np_1m          — 2D array of 1-minute candle data.
        index_datetime — Column index for datetime.
        index_open     — Column index for open price.
        index_high     — Column index for high price.
        index_low      — Column index for low price.
    """
    n_minutes = n_days * 24 * 60  # Total 1-minute candles

    # Generate datetime index at 1-minute frequency
    datetimes = pd.date_range(start='2024-01-01', periods=n_minutes, freq='1min').values

    # Simulate price via geometric random walk
    np.random.seed(42)
    returns      = np.random.normal(0, 0.001, n_minutes)  # Small random returns
    close_prices = base_price * np.cumprod(1 + returns)

    # Construct OHLCV columns
    open_prices  = np.roll(close_prices, 1);  open_prices[0] = base_price
    high_prices  = np.maximum(open_prices, close_prices) * (1 + np.abs(np.random.normal(0, 0.0005, n_minutes)))
    low_prices   = np.minimum(open_prices, close_prices) * (1 - np.abs(np.random.normal(0, 0.0005, n_minutes)))
    volumes      = np.random.randint(1, 100, n_minutes).astype(float)

    # Stack into structured array: [datetime, open, high, low, close, volume]
    # Convert datetimes to object type to allow mixed types in np.column_stack
    np_1m = np.column_stack([
        datetimes.astype(object), # Convert to object to allow mixed types
        open_prices,
        high_prices,
        low_prices,
        close_prices,
        volumes
    ])

    # Fixed column indices matching dataset column order
    index_datetime = 0  # Column 0: Datetime
    index_open     = 1  # Column 1: Open
    index_high     = 2  # Column 2: High
    index_low      = 3  # Column 3: Low

    return np_1m, index_datetime, index_open, index_high, index_low


# The actual data generation and instantiation will be moved to the main run function.
# print(f"1-minute data shape : {obj_dataset.np_1m.shape}")
# print(f"Datetime range      : {raw_1m_data[0, 0]}  →  {raw_1m_data[-1, 0]}")

def generate_mock_predictions(n_periods: int = 60, freq: str = '4h') -> pd.DataFrame:
    """
    Generate a synthetic prediction signal DataFrame.

    Parameters
    ----------
    n_periods : int
        Number of prediction intervals.
    freq : str
        Frequency of each prediction candle (pandas offset alias).

    Returns
    -------
    pd.DataFrame
        Columns: ['datetime', 'predicted_direction']
    """
    datetimes  = pd.date_range(start='2024-01-01', periods=n_periods, freq=freq)

    # Randomly assign long (+1), short (-1), or neutral (0)
    np.random.seed(99)
    directions = np.random.choice([1, -1, 0], size=n_periods, p=[0.45, 0.45, 0.10])

    return pd.DataFrame({'datetime': datetimes, 'predicted_direction': directions})

# The actual prediction generation will be moved to the main run function.
# df_predictions = generate_mock_predictions(n_periods=55, freq='4h')

# print(f"Prediction periods : {len(df_predictions)}")
# print(f"Signal distribution:\n{df_predictions['predicted_direction'].value_counts().to_string()}")
# df_predictions.head(10)

# The `Backtest` class definition has been moved to the `imports-cell` to ensure it's defined before usage.

def run_backtest_simulation(
    n_days_data: int = 10,
    base_price_data: float = 30000.0,
    n_periods_predictions: int = 55,
    freq_predictions: str = '4h',
    starting_balance: float = 1000,
    take_profit: float = 1.0,
    stop_loss: float = 1.0,
    buy_after_minutes: int = 0,
    transaction_fee: float = 0.05,
    leverage: float = 1.0,
    slippage: float = 0.0,
    output_path: str = "backtest_ledger_vectorized.csv"
):
    """
    Orchestrates the entire vectorized backtesting simulation using functions.

    Parameters
    ----------
    n_days_data : int
        Number of trading days to simulate for 1-minute OHLCV data.
    base_price_data : float
        Starting price for the random walk simulation of 1-minute data.
    n_periods_predictions : int
        Number of prediction intervals to generate.
    freq_predictions : str
        Frequency of each prediction candle (pandas offset alias).
    starting_balance : float
        Initial account capital in USD.
    take_profit : float
        Take-profit threshold as a percentage (e.0g., 1.0 = 1%).
    stop_loss : float
        Stop-loss threshold as a percentage.
    buy_after_minutes : int
        Number of 1-minute candles to skip before executing entry.
    transaction_fee : float
        Per-trade transaction fee as a percentage of position size.
    leverage : float
        Leverage multiplier applied to PnL calculations.
    slippage : float
        Simulated price slippage per trade as a percentage.
    output_path : str
        File path to save the trade ledger CSV.
    """

    print("--- Generating Mock Data ---")
    np_1m, index_1m_datetime, index_1m_open, index_1m_high, index_1m_low = generate_mock_1m_data(n_days=n_days_data, base_price=base_price_data)
    print(f"  1-minute data shape : {np_1m.shape}")
    print(f"  Datetime range      : {np_1m[0, index_1m_datetime]}  →  {np_1m[-1, index_1m_datetime]}")

    print("\n--- Generating Mock Predictions ---")
    df_predictions = generate_mock_predictions(n_periods=n_periods_predictions, freq=freq_predictions)
    # Convert predictions to numpy for faster access
    df_predictions["datetime"] = pd.to_datetime(df_predictions["datetime"])
    np_model_predictions  = df_predictions.to_numpy()
    index_pred_datetime  = df_predictions.columns.get_loc("datetime")
    index_pred_direction = df_predictions.columns.get_loc("predicted_direction")

    print(f"  Prediction periods : {len(df_predictions)}")
    print(f"  Signal distribution:\n{df_predictions['predicted_direction'].value_counts().to_string()}")

    print("\n--- Running Backtest Simulation ---")

    # --- Initialize state variables ---
    current_balance = starting_balance
    breaking_balance = starting_balance * 0.5  # 50% drawdown hard stop
    buy_price = 0
    sell_price = 0
    in_position = False
    array_to_save = []

    # --- Risk parameters ---
    take_profit_percent = take_profit / 100
    stop_loss_percent = stop_loss / 100
    transaction_fee_percent = transaction_fee * leverage

    # Initialize direction tracking from the first prediction signal
    previous_pred_direction = np_model_predictions[0][index_pred_direction]
    current_pred_direction = previous_pred_direction
    break_on_huge_loss = False

    # --- Main iteration: process each prediction interval ---
    for i in range(0, len(np_model_predictions) - 1):
        current_pred_direction = np_model_predictions[i][index_pred_direction]

        # ── Neutral signal handling ─────────────────────────────────────
        if current_pred_direction == 0:
            # If in position, record a passive interval marker and check TP/SL
            if in_position:
                np_temp, np_temp_high, np_temp_low = get_interval_min_data(
                    np_1m, np_model_predictions, i, index_1m_datetime, index_1m_high, index_1m_low, index_pred_datetime
                )
                if len(np_temp) > 10:
                    record_trade(
                        array_to_save, np_temp[buy_after_minutes][index_1m_datetime],
                        current_pred_direction, "same direction", buy_price, sell_price, current_balance, 0
                    )
                    # Check TP/SL for current position (if any)
                    if in_position:
                        if previous_pred_direction > 0:  # LONG
                            tp_amount = buy_price * (1 + take_profit_percent)
                            sl_amount   = buy_price * (1 - stop_loss_percent)
                        else:                                  # SHORT
                            tp_amount = buy_price * (1 - take_profit_percent)
                            sl_amount   = buy_price * (1 + stop_loss_percent)

                        triggered, df_temp_index, current_sell_price = find_tp_sl_index(
                            current_pred_direction, tp_amount, sl_amount, np_temp_high, np_temp_low
                        )
                        if triggered:
                            pnl = calculate_pnl(
                                buy_price, current_sell_price, previous_pred_direction,
                                leverage, transaction_fee_percent, slippage
                            )
                            current_balance += current_balance * (pnl / 100)
                            in_position = False
                            exit_label = " - take_profit" if pnl > 0 else " - stop_loss"
                            record_trade(
                                array_to_save, np_temp[df_temp_index][index_1m_datetime],
                                previous_pred_direction, "sell" + exit_label, buy_price, current_sell_price, current_balance, pnl
                            )
            continue  # No new position opened on neutral signal

        # ── Fetch 1-minute data for this interval ───────────────────────
        np_temp, np_temp_high, np_temp_low = get_interval_min_data(
            np_1m, np_model_predictions, i, index_1m_datetime, index_1m_high, index_1m_low, index_pred_datetime
        )
        if not len(np_temp) > 10:
            continue  # Insufficient data — skip interval

        # ── Same direction: record passive marker (no new entry) ─────────
        if in_position and previous_pred_direction == current_pred_direction:
            record_trade(
                array_to_save, np_temp[buy_after_minutes][index_1m_datetime],
                current_pred_direction, "same direction", buy_price, sell_price, current_balance, 0
            )

        # ── New position entry ───────────────────────────────────────────
        if not in_position:
            buy_price = np_temp[buy_after_minutes][index_1m_open]
            pnl = transaction_fee_percent * -1
            pnl -= slippage
            current_balance += current_balance * (pnl / 100)
            in_position = True
            record_trade(
                array_to_save, np_temp[buy_after_minutes][index_1m_datetime],
                current_pred_direction, "buy", buy_price, 0, current_balance, pnl
            )

        # ── Direction reversal: close position and re-enter ──────────────
        if current_pred_direction != previous_pred_direction:
            if len(np_temp) >= 10:
                sell_price  = np_temp[buy_after_minutes][index_1m_open]
                sell_datetime    = np_temp[buy_after_minutes][index_1m_datetime]
                pnl = calculate_pnl(
                    buy_price, sell_price, previous_pred_direction,
                    leverage, transaction_fee_percent, slippage
                )
                current_balance += current_balance * (pnl / 100)
                in_position = False
                record_trade(array_to_save, sell_datetime, previous_pred_direction, "sell - direction change", buy_price, sell_price, current_balance, pnl)

                # Immediately re-enter in the new direction
                buy_price = np_temp[buy_after_minutes][index_1m_open]
                pnl_entry = transaction_fee_percent * -1 # Re-entry fee
                pnl_entry -= slippage # Re-entry slippage
                current_balance += current_balance * (pnl_entry / 100)
                in_position = True
                record_trade(
                    array_to_save, np_temp[buy_after_minutes][index_1m_datetime],
                    current_pred_direction, "buy", buy_price, 0, current_balance, pnl_entry
                )
            else:
                continue

        # ── TP/SL check at 1-minute resolution ──────────────────────────
        if in_position:
            if current_pred_direction > 0:  # LONG
                tp_amount = buy_price * (1 + take_profit_percent)
                sl_amount   = buy_price * (1 - stop_loss_percent)
            else:                                  # SHORT
                tp_amount = buy_price * (1 - take_profit_percent)
                sl_amount   = buy_price * (1 + stop_loss_percent)

            triggered, df_temp_index, current_sell_price = find_tp_sl_index(
                current_pred_direction, tp_amount, sl_amount, np_temp_high, np_temp_low
            )

            if triggered:
                pnl = calculate_pnl(
                    buy_price, current_sell_price, current_pred_direction,
                    leverage, transaction_fee_percent, slippage
                )
                current_balance += current_balance * (pnl / 100)
                in_position = False
                exit_label = " - take_profit" if pnl > 0 else " - stop_loss"
                record_trade(
                    array_to_save, np_temp[df_temp_index][index_1m_datetime],
                    current_pred_direction, "sell" + exit_label, buy_price, current_sell_price, current_balance, pnl
                )

        previous_pred_direction = current_pred_direction

        # ── Maximum drawdown hard stop ───────────────────────────────────
        if current_balance < breaking_balance:
            break_on_huge_loss = True
            break

    # --- Construct trade ledger DataFrame ---
    df_ledger = pd.DataFrame(array_to_save, columns=[
        "datetime",
        "predicted_direction",
        "action",
        "buy_price",
        "sell_price",
        "balance",
        "pnl",
    ])
    df_ledger["pnl_sum"] = df_ledger["pnl"].cumsum()  # Cumulative PnL
    df_ledger[["balance", "pnl", "pnl_sum"]] = df_ledger[["balance", "pnl", "pnl_sum"]].round(2)

    # --- Return results ---
    if len(df_ledger) > 1:
        pnl_percent = np.round(df_ledger["pnl_sum"].iloc[-1], 2)
        if break_on_huge_loss:
            final_balance = -1000 # Indicates hard stop triggered
        else:
            final_balance = round(current_balance, 2)
    else:
        final_balance = 0
        pnl_percent = 0

    # --- Summary output ---
    print("\n" + "=" * 45)
    print(f"  Starting Balance : ${starting_balance:,.2f}")
    print(f"  Final Balance    : ${final_balance:,.2f}")
    print(f"  Total PnL (%)   : {pnl_percent:.2f}%%")
    print(f"  Total Trades     : {len(df_ledger)}")
    print("=" * 45)

    # --- Inspect Trade Ledger ---
    print("\n--- Trade Ledger ---")
    pd.set_option('display.max_rows', 80)
    pd.set_option('display.float_format', '{:.4f}'.format)
    print(f"Total records: {len(df_ledger)}")
    display(df_ledger)

    # --- Performance Analysis ---
    print("\n--- Performance Analysis ---")
    df_sells = df_ledger[df_ledger['action'].str.startswith('sell')].copy()

    if len(df_sells) > 0:
        winning_trades = df_sells[df_sells['pnl'] > 0]
        losing_trades  = df_sells[df_sells['pnl'] <= 0]

        win_rate     = len(winning_trades) / len(df_sells) * 100 if len(df_sells) > 0 else 0
        avg_win_pnl  = winning_trades['pnl'].mean() if len(winning_trades) > 0 else 0
        avg_loss_pnl = losing_trades['pnl'].mean()  if len(losing_trades)  > 0 else 0
        profit_factor = (
            winning_trades['pnl'].sum() / abs(losing_trades['pnl'].sum())
            if len(losing_trades) > 0 and losing_trades['pnl'].sum() != 0 else float('inf')
        )

        print("=" * 45)
        print("  PERFORMANCE METRICS")
        print("=" * 45)
        print(f"  Total Exit Events    : {len(df_sells)}")
        print(f"  Winning Trades       : {len(winning_trades)}")
        print(f"  Losing Trades        : {len(losing_trades)}")
        print(f"  Win Rate             : {win_rate:.1f}%%")
        print(f"  Avg Win PnL          : {avg_win_pnl:.4f}%%")
        print(f"  Avg Loss PnL         : {avg_loss_pnl:.4f}%%")
        print(f"  Profit Factor        : {profit_factor:.2f}")
        print(f"  Cumulative PnL       : {pnl_percent:.2f}%%")
        print(f"  Final Balance        : ${final_balance:,.2f}")
        print("=" * 45)

        # Action distribution
        print("\nAction Distribution:")
        print(df_ledger['action'].value_counts().to_string())
    else:
        print("No sell events recorded.")

    # --- Export Ledger to CSV ---
    # print(f"\n--- Exporting Ledger ---")
    # df_ledger.to_csv(output_path, index=False)
    # print(f"Ledger exported to: {output_path}")

    return df_ledger, final_balance, pnl_percent

# --- Call the main backtest simulation function ---
# You can customize parameters here
result_ledger, result_final_balance, result_pnl_percent = run_backtest_simulation(
    starting_balance = 1000,
    take_profit      = 1.0,
    stop_loss        = 1.0,
    buy_after_minutes = 0,
    transaction_fee  = 0.05,
    leverage         = 1.0,
    slippage         = 0.0
)