# -*- coding: utf-8 -*-
"""Version 2 VIX-Filtered Mean Reversion Futures Strategy.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1hVHcFDDox4gpxfu6OR6V8N1ZcMCA0Io3
"""

# @title
# Install required packages
!pip install -q backtesting yfinance pandas numpy
!pip install -q scikit-optimize
!pip install -q sambo
!pip install -q bokeh>=2.4.0
!pip install scikit-optimize
print("✓ All packages installed successfully!")

# @title
import warnings
warnings.filterwarnings('ignore')

import yfinance as yf
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from backtesting import Backtest, Strategy
from backtesting.lib import crossover
import logging

logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
logger = logging.getLogger(__name__)

print("✓ All libraries imported successfully!")

# @title
def fetch_futures_data(ticker='ES', max_days=729):
    """
    Fetch futures data from Yahoo Finance (1h bars)
    Auto-adjusts to last 729 days (Yahoo limit)
    """
    FUTURES_TICKERS = {
        'ES': 'ES=F',
        'NQ': 'NQ=F',
        'MES': 'MES=F',
        'MNQ': 'MNQ=F'
    }

    yahoo_ticker = FUTURES_TICKERS.get(ticker, 'ES=F')

    # Calculate date range
    end_date = datetime.now()
    start_date = end_date - timedelta(days=max_days)

    logger.info(f"Fetching {ticker} data from {start_date.date()} to {end_date.date()}")

    try:
        ticker_obj = yf.Ticker(yahoo_ticker)
        data = ticker_obj.history(
            start=start_date.strftime('%Y-%m-%d'),
            end=end_date.strftime('%Y-%m-%d'),
            interval='1h',
            auto_adjust=False
        )

        if data.empty:
            logger.error(f"No data retrieved for {ticker}")
            return pd.DataFrame()

        # Prepare columns
        data.columns = [col.title() for col in data.columns]
        data = data[['Open', 'High', 'Low', 'Close', 'Volume']]
        data = data.dropna()

        # Timezone handling
        if data.index.tz is None:
            data.index = data.index.tz_localize('America/New_York', ambiguous='infer', nonexistent='shift_forward')
        else:
            data.index = data.index.tz_convert('America/New_York')

        logger.info(f"✓ Fetched {len(data)} bars for {ticker}")
        return data

    except Exception as e:
        logger.error(f"Error fetching {ticker}: {str(e)}")
        return pd.DataFrame()


def fetch_vix_data(max_days=729):
    """Fetch VIX data for regime filtering"""

    end_date = datetime.now()
    start_date = end_date - timedelta(days=max_days)

    logger.info(f"Fetching VIX data")

    try:
        vix = yf.Ticker('^VIX')
        data = vix.history(
            start=start_date.strftime('%Y-%m-%d'),
            end=end_date.strftime('%Y-%m-%d'),
            interval='1d',
            auto_adjust=False
        )

        if data.empty:
            logger.error("No VIX data retrieved")
            return pd.DataFrame()

        data = data[['Close']].rename(columns={'Close': 'VIX'})
        logger.info(f"✓ Fetched {len(data)} VIX data points")
        return data

    except Exception as e:
        logger.error(f"Error fetching VIX: {str(e)}")
        return pd.DataFrame()


def merge_data(futures_data, vix_data):
    """Merge VIX with futures data"""

    if futures_data.empty or vix_data.empty:
        logger.error("Cannot merge: Empty data")
        return pd.DataFrame()

    # Forward fill VIX to match hourly futures data
    vix_reindexed = vix_data.reindex(futures_data.index, method='ffill')

    merged = futures_data.copy()
    merged['VIX'] = vix_reindexed['VIX']
    merged = merged.dropna(subset=['VIX'])

    logger.info(f"✓ Merged  {len(merged)} bars")
    return merged

print("✓ Data fetcher functions defined!")

# @title
# Indicators

import pandas as pd
import numpy as np


def calculate_vwap(close, volume, high, low, index):
    """
    Calculate VWAP (Volume Weighted Average Price) - resets daily

    Formula: VWAP = Σ(Typical Price × Volume) / Σ(Volume)
    Where Typical Price = (High + Low + Close) / 3

    The calculation resets at the start of each trading day.
    """
    typical_price = (high + low + close) / 3

    df = pd.DataFrame({
        'tp': typical_price,
        'vol': volume,
        'date': index.date
    })

    # Calculate cumulative sums grouped by date
    df['cum_tp_vol'] = df.groupby('date')['tp'].transform(
        lambda x: (x * df.loc[x.index, 'vol']).cumsum()
    )
    df['cum_vol'] = df.groupby('date')['vol'].transform('cumsum')

    # VWAP = Cumulative(TP × Vol) / Cumulative(Vol)
    vwap = df['cum_tp_vol'] / df['cum_vol']

    return vwap.values


def calculate_vwap_bands(close, volume, high, low, index, std_mult=1.0):
    """
    Calculate VWAP Standard Deviation Bands using UNWEIGHTED standard deviation

    Formula:
    1. Calculate VWAP first
    2. For each bar: deviation = (Typical Price - VWAP)
    3. Variance = Σ(deviation²) / N  (where N = number of bars)
    4. Std Dev = sqrt(Variance)
    5. Upper Band = VWAP + (std_mult × Std Dev)
    6. Lower Band = VWAP - (std_mult × Std Dev)

    This matches ThinkorSwim and standard statistical definition.
    Each bar contributes equally to the standard deviation regardless of volume.
    """
    # First calculate VWAP
    vwap = calculate_vwap(close, volume, high, low, index)
    typical_price = (high + low + close) / 3

    df = pd.DataFrame({
        'tp': typical_price,
        'vwap': vwap,
        'date': index.date
    })

    # Calculate squared deviations from VWAP
    df['sq_diff'] = (df['tp'] - df['vwap']) ** 2

    # Calculate unweighted variance (sum of squared differences / count)
    df['cum_sq_diff'] = df.groupby('date')['sq_diff'].transform('cumsum')
    df['count'] = df.groupby('date').cumcount() + 1  # Running count of bars

    # Variance = Cumulative Sum of Squared Differences / Count of Bars
    variance = df['cum_sq_diff'] / df['count']
    std_dev = np.sqrt(variance)

    # Calculate bands
    upper = vwap + (std_mult * std_dev)
    lower = vwap - (std_mult * std_dev)

    return upper, lower


def calculate_vwap_slope(close, volume, high, low, index, lookback=6):
    """
    Calculate VWAP slope as percentage change

    Formula: Slope = ((VWAP_current - VWAP_lookback) / VWAP_lookback) × 100

    Args:
        lookback: Number of periods to look back for slope calculation

    Returns:
        Slope as percentage
    """
    vwap = calculate_vwap(close, volume, high, low, index)
    slope = pd.Series(vwap).pct_change(periods=lookback) * 100

    return slope.values


def calculate_atr(high, low, close, period=14):
    """
    Calculate ATR (Average True Range)

    True Range is the maximum of:
    1. Current High - Current Low
    2. |Current High - Previous Close|
    3. |Current Low - Previous Close|

    ATR is the Exponential Moving Average of True Range over the period.
    """
    h = pd.Series(high)
    l = pd.Series(low)
    c = pd.Series(close)

    # Calculate the three components of True Range
    tr1 = h - l
    tr2 = abs(h - c.shift())
    tr3 = abs(l - c.shift())

    # True Range is the maximum of the three
    tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)

    # Apply exponential moving average
    atr = tr.ewm(span=period, adjust=False).mean()

    return atr.values


def calculate_volume_ma(volume, period=9):
    """
    Calculate Volume Moving Average

    Simple moving average of volume over the specified period.
    """
    return pd.Series(volume).rolling(window=period).mean().values


def calculate_volume_profile(high, low, close, volume, index, num_bins=24):
    """
    Calculate Volume Profile metrics: VAH (Value Area High), VAL (Value Area Low), POC (Point of Control)

    Process:
    1. Divide the day's price range into bins (default 24)
    2. Distribute volume across bins based on where price traded
    3. POC = price level with highest volume
    4. Value Area = 70% of total volume centered around POC
    5. VAH = upper boundary of value area
    6. VAL = lower boundary of value area

    The calculation resets daily.
    """
    df = pd.DataFrame({
        'high': high,
        'low': low,
        'close': close,
        'volume': volume,
        'date': index.date
    })

    vah_list, val_list, poc_list = [], [], []

    for date in df['date'].unique():
        day_data = df[df['date'] == date]

        # Handle insufficient data
        if len(day_data) < 3:
            vah_list.extend([np.nan] * len(day_data))
            val_list.extend([np.nan] * len(day_data))
            poc_list.extend([np.nan] * len(day_data))
            continue

        price_min = day_data['low'].min()
        price_max = day_data['high'].max()

        # Handle no price movement
        if price_max == price_min:
            vah_list.extend([day_data['close'].iloc[0]] * len(day_data))
            val_list.extend([day_data['close'].iloc[0]] * len(day_data))
            poc_list.extend([day_data['close'].iloc[0]] * len(day_data))
            continue

        # Create price bins
        bins = np.linspace(price_min, price_max, num_bins)
        vol_profile = np.zeros(len(bins) - 1)

        # Distribute volume across bins based on bar overlap
        for _, row in day_data.iterrows():
            for i in range(len(bins) - 1):
                if row['low'] <= bins[i+1] and row['high'] >= bins[i]:
                    # Calculate overlap between bar range and bin range
                    overlap = min(row['high'], bins[i+1]) - max(row['low'], bins[i])
                    bar_range = row['high'] - row['low']

                    if bar_range > 0:
                        # Distribute volume proportionally to overlap
                        vol_profile[i] += row['volume'] * (overlap / bar_range)
                    else:
                        # Equal distribution if no range
                        vol_profile[i] += row['volume'] / (len(bins) - 1)

        # Find POC (Point of Control) - bin with highest volume
        poc_idx = np.argmax(vol_profile)
        poc_price = (bins[poc_idx] + bins[poc_idx + 1]) / 2

        total_vol = vol_profile.sum()
        if total_vol == 0:
            vah_list.extend([day_data['close'].mean()] * len(day_data))
            val_list.extend([day_data['close'].mean()] * len(day_data))
            poc_list.extend([day_data['close'].mean()] * len(day_data))
            continue

        # Calculate Value Area (70% of volume around POC)
        va_vol = total_vol * 0.70
        va_indices = [poc_idx]
        current_vol = vol_profile[poc_idx]

        lower_idx = poc_idx - 1
        upper_idx = poc_idx + 1

        # Expand from POC until 70% volume is captured
        while current_vol < va_vol:
            lower_vol = vol_profile[lower_idx] if lower_idx >= 0 else 0
            upper_vol = vol_profile[upper_idx] if upper_idx < len(vol_profile) else 0

            if lower_vol >= upper_vol and lower_idx >= 0:
                va_indices.append(lower_idx)
                current_vol += lower_vol
                lower_idx -= 1
            elif upper_idx < len(vol_profile):
                va_indices.append(upper_idx)
                current_vol += upper_vol
                upper_idx += 1
            else:
                break

        # Calculate VAH and VAL
        vah = bins[max(va_indices) + 1] if max(va_indices) + 1 < len(bins) else bins[-1]
        val = bins[min(va_indices)]

        vah_list.extend([vah] * len(day_data))
        val_list.extend([val] * len(day_data))
        poc_list.extend([poc_price] * len(day_data))

    return np.array(vah_list), np.array(val_list), np.array(poc_list)


print("✓ All indicators defined correctly!")

# @title
# Mean Reversion Futures Strategy - REVISED
# Professional Implementation with backtesting.py

import pandas as pd
import numpy as np


class MeanReversionStrategy(Strategy):
    """
    Mean Reversion Futures Strategy - 1 Hour Bars

    Entry Logic:
    - VIX regime filtering (low/mid volatility)
    - Price touches VWAP bands (1σ or 2σ)

    Exit Logic:
    - Stop Loss: 2.0x ATR
    - Take Profit: 100% at VWAP touch
    - Profit Target: $100 USD per contract

    Position Management:
    - 1 contract per trade
    - Maximum 3 concurrent positions
    """

    # ========================
    # STRATEGY PARAMETERS
    # ========================

    # VIX Regime Thresholds
    vix_low_threshold = 12      # Below this: use 1σ and 2σ bands
    vix_mid_threshold = 27      # Above this: no entries

    # VWAP Bands
    bb_std_1 = 1.0             # First standard deviation multiplier
    bb_std_2 = 2.0             # Second standard deviation multiplier

    # ATR Settings
    atr_period = 14            # ATR calculation period
    stop_multiplier = 3.0      # Stop loss distance (ATR multiplier)

    # Profit Target (DISABLED)
    profit_target_usd = 100.0  # Take profit at $100 per contract

    # Volume Indicators (DISABLED)
    volume_ma_period = 9       # Volume MA period (not used)

    # VWAP Slope (DISABLED)
    vwap_slope_lookback = 9    # VWAP slope lookback (not used)
    vwap_slope_threshold = 0.15  # VWAP slope threshold (not used)

    # Position Management
    max_positions = 3          # Maximum concurrent positions
    # max_trades_per_day = 3   # DISABLED: No daily trade limit

    # Session Management (DISABLED)
    # session_start_hour = 10  # DISABLED: Trade all hours

    # Risk Management (DISABLED)
    # daily_loss_limit = 0.03  # DISABLED: No daily loss limit

    # Contract Specifications
    price_per_point = 2.0     # ES=50, MES=5, MNQ=2


    def init(self):
        """Initialize indicators and tracking variables"""

        # ========================
        # REQUIRED DATA VALIDATION
        # ========================
        if 'VIX' not in self.data.df.columns:
            raise ValueError("VIX data not found in DataFrame")

        # ========================
        # CORE INDICATORS
        # ========================

        # VWAP (Volume Weighted Average Price)
        self.vwap = self.I(
            calculate_vwap,
            self.data.Close, self.data.Volume,
            self.data.High, self.data.Low,
            self.data.index,
            name='VWAP'
        )

        # VWAP Bands - 1 Standard Deviation
        bb_result_1 = self.I(
            calculate_vwap_bands,
            self.data.Close, self.data.Volume,
            self.data.High, self.data.Low,
            self.data.index,
            self.bb_std_1,
            name='BB1'
        )
        self.bb_upper_1 = bb_result_1[0] # ✓ +1σ band (VWAP + 1×StdDev)
        self.bb_lower_1 = bb_result_1[1] # ✓ -1σ band (VWAP - 1×StdDev)

        # VWAP Bands - 2 Standard Deviations
        bb_result_2 = self.I(
            calculate_vwap_bands,
            self.data.Close, self.data.Volume,
            self.data.High, self.data.Low,
            self.data.index,
            self.bb_std_2,
            name='BB2'
        )
        self.bb_upper_2 = bb_result_2[0] # ✓ +2σ band (VWAP + 2×StdDev)
        self.bb_lower_2 = bb_result_2[1] # ✓ -2σ band (VWAP - 2×StdDev)

        # ATR (Average True Range) for stop loss
        self.atr = self.I(
            calculate_atr,
            self.data.High, self.data.Low, self.data.Close,
            self.atr_period,
            name='ATR'
        )

        # ========================
        # VIX DATA
        # ========================
        self.vix = self.data.df['VIX'].values

        # ========================
        # TRACKING VARIABLES
        # ========================
        self.current_bar = 0


    def next(self):
        """Main strategy logic executed on each bar"""

        self.current_bar += 1

        # ========================
        # DATA VALIDATION
        # ========================
        if np.isnan(self.vwap[-1]) or np.isnan(self.atr[-1]):
            return

        # ========================
        # MANAGE EXISTING POSITIONS
        # ========================
        open_trades = getattr(self, 'trades_open', [])
        if open_trades:
            self._manage_exits()

        # ========================
        # POSITION LIMIT CHECK
        # ========================
        if len(open_trades) >= self.max_positions:
            return  # Maximum 3 concurrent positions

        # ========================
        # CURRENT MARKET DATA
        # ========================
        vix_current = self.vix[-1]
        close_current = self.data.Close[-1]

        # ========================
        # ENTRY LOGIC
        # ========================

        # Filter 1: VIX Regime
        if vix_current > self.vix_mid_threshold:
            return  # No entries in high volatility

        # Determine which bands are valid for entry based on VIX
        if vix_current < self.vix_low_threshold:
            # Low VIX: can enter at 1σ or 2σ
            entry_band_1 = self.bb_lower_1[-1]
            entry_band_2 = self.bb_lower_2[-1]
        else:
            # Mid VIX (12-27): only enter at 2σ
            entry_band_1 = None
            entry_band_2 = self.bb_lower_2[-1]

        # Check if price is at qualifying band
        price_at_band_1 = entry_band_1 is not None and close_current <= entry_band_1
        price_at_band_2 = close_current <= entry_band_2

        if not (price_at_band_1 or price_at_band_2):
            return  # Price not at any qualifying band

        # ========================
        # EXECUTE ENTRY
        # ========================

        # Calculate stop loss at entry
        stop_loss_entry = close_current - (self.atr[-1] * self.stop_multiplier)

        # Enter with 1 contract
        trade = self.buy(size=1, stop=stop_loss_entry)

        if trade:
            # Store entry details on trade object
            trade.entry_price = close_current
            trade.entry_bar = self.current_bar

            # Calculate profit target in price points
            profit_target_points = self.profit_target_usd / self.price_per_point
            trade.profit_target_price = close_current + profit_target_points


    def _manage_exits(self):
        """
        Exit management for all open positions

        Exit Rules:
        1. Stop Loss: 2.0x ATR below entry
        2. VWAP Touch: Close 100% at VWAP
        3. Profit Target: Close 100% at $100 profit
        """

        close_current = self.data.Close[-1]
        vwap_current = self.vwap[-1]

        # Process each open trade
        open_trades = list(getattr(self, 'trades_open', []))

        for trade in open_trades:
            # Ensure trade attributes exist
            if not hasattr(trade, 'entry_price'):
                trade.entry_price = trade.entry_price
            if not hasattr(trade, 'profit_target_price'):
                # Calculate if missing
                profit_target_points = self.profit_target_usd / self.price_per_point
                trade.profit_target_price = trade.entry_price + profit_target_points

            # ========================
            # EXIT 1: ATR
            # ========================
            stop_loss = trade.entry_price - (self.atr[-1] * self.stop_multiplier)
            if close_current <= stop_loss:
                trade.close()
                continue

            # ========================
            # EXIT 2: UPPER BAND TARGET (VIX-Based)
            # ========================

            # Determine which upper bands are valid for exit based on VIX
#            if vix_current < self.vix_low_threshold:
#             # Low VIX: can exit at +1σ or +2σ
#              exit_band_1 = self.bb_upper_1[-1]
#                exit_band_2 = self.bb_upper_2[-1]

            # Exit if price touches either band
#            if close_current >= exit_band_1 or close_current >= exit_band_2:
#                trade.close()  # Close 100% of position
#                continue

#            else:
            # Mid/High VIX: only exit at +2σ (more conservative, let winners run)
#                exit_band_2 = self.bb_upper_2[-1]

#            if close_current >= exit_band_2:
#                trade.close()  # Close 100% of position
#                continue

            # ========================
            # EXIT 2.2: VWAP TOUCH (100% Exit)
            # ========================
            #if close_current >= vwap_current:
                #trade.close()  # Close 100% of position
                #continue

            # ========================
            # EXIT 3: PROFIT TARGET ($100)
            # ========================
            if close_current >= trade.profit_target_price:
                trade.close()  # Close 100% of position
                continue


print("✓ MeanReversionStrategy class defined successfully!")

# @title
# Download data
print("="*60)
print("DOWNLOADING DATA")
print("="*60)

# Fetch ES futures
es_data = fetch_futures_data('ES', max_days=729)

if es_data.empty:
    raise ValueError("Failed to fetch ES data")

# Fetch VIX
vix_data = fetch_vix_data(max_days=729)

if vix_data.empty:
    raise ValueError("Failed to fetch VIX data")

# Merge
data = merge_data(es_data, vix_data)

if data.empty:
    raise ValueError("Failed to merge data")

print("\n" + "="*60)
print("DATA READY")
print("="*60)
print(f"Total bars: {len(data):,}")
print(f"Date range: {data.index[0]} to {data.index[-1]}")
print(f"Columns: {list(data.columns)}")
print("\nFirst 3 rows:")
print(data.head(3))

# @title
# Backtest
import pandas as pd # Ensure pandas is imported for set_option

# Define custom commission function for a flat fee per order
def custom_commission(size, price):
    # $0.25 flat fee per order (buy or sell)
    return 0.25

bt = Backtest(
    data,
    MeanReversionStrategy,
    cash=30000,
    commission=custom_commission, # Use the custom callable for commission
    margin=0.05,
    trade_on_close=False,
    hedging=False,
    exclusive_orders=True
)

# Run backtest
stats = bt.run()

#Results table
print("\n" + "="*60)
print("BACKTEST RESULTS")
print("="*60)
print(stats)

# Trades
# Export trade log
trades = stats['_trades']

# Columns to drop
columns_to_drop = [
    'Entry_Hammer', 'Exit_Hammer',
    'Entry_Engulfing', 'Exit_Engulfing',
    'Entry_Doji', 'Exit_Doji'
]

# Drop columns if they exist
trades = trades.drop(columns=[col for col in columns_to_drop if col in trades.columns])

print("\n" + "="*60)
print(f"Total trades executed: {len(trades)}")
print("="*60)
print("Trade log (all trades):")

# Set display options to show all columns and prevent wrapping
pd.set_option('display.max_columns', None)
pd.set_option('display.width', 1000) # Adjust as needed for your screen
pd.set_option('display.max_rows', None) # Show all rows

#print(trades) Disabled for now

# Reset display options to default after printing (optional, but good practice)
pd.reset_option('display.max_columns')
pd.reset_option('display.width')
pd.reset_option('display.max_rows')

# Save to CSV
trades.to_csv('trades.csv')
print("\n Trades exported to trades.csv")

# @title
stats.tail()

# @title
stats['_equity_curve']

# @title
stats['_trades']

# @title
# Native backtesting.py plot
bt.plot()

# @title
# Optimizer
stats_optimized, heatmap, optimize_result = bt.optimize(
        vix_low_threshold=[10, 15],
        vix_mid_threshold=[20, 25, 30],
        atr_period=[10, 14, 18],
        risk_per_trade=[0.005, 0.01, 0.015],
        stop_multiplier=[1.0, 1.5, 2.0],
        first_target_portion=[0.5, 0.7, 0.8],
        constraint=lambda p: p.vix_low_threshold < p.vix_mid_threshold,
        maximize='Equity Final [$]', #Change to % if needed or sharpe ratio
        method='sambo',
        max_tries=100,
        random_state=42,
        return_heatmap=True,
        return_optimization=True
    )
# Results
print("\n" + "="*60)
print("OPTIMIZATION VERIFICATION")
print("="*60)

# Check the stats object
if stats_optimized is not None:
    print("Optimization completed successfully")
    print(f"   Optimal strategy: {stats_optimized._strategy}")
    print(f"   Total trades in optimized backtest: {stats_optimized['# Trades']}")
    print("\n--- Full Optimized Strategy Statistics ---")
    print(stats_optimized) # Print the entire stats_optimized object
    print("----------------------------------------")
else:
    print("Optimization failed")

# @title
print("\n" + "="*60)
print("COMPARISON: DEFAULT vs OPTIMIZED")
print("="*60)

if stats_optimized is not None:
    comparison = pd.DataFrame({
        'Metric': ['Return [%]', 'Sharpe Ratio', 'Max Drawdown [%]', 'Win Rate [%]', '# Trades'],
        'Default': [
            stats['Return [%]'],
            stats['Sharpe Ratio'],
            stats['Max. Drawdown [%]'],
            stats['Win Rate [%]'],
            stats['# Trades']
        ],
        'Optimized': [
            stats_optimized['Return [%]'],
            stats_optimized['Sharpe Ratio'],
            stats_optimized['Max. Drawdown [%]'],
            stats_optimized['Win Rate [%]'],
            stats_optimized['# Trades']
        ]
    })

    comparison['Improvement %'] = ((comparison['Optimized'] - comparison['Default']) / comparison['Default'] * 100).round(2)

    print(comparison.to_string(index=False))

    print("\n✓ Analysis complete!")
else:
    print("Optimization did not complete successfully, cannot perform comparison.")

# @title
# Best combination
print("\n" + "="*60)
print("TOP 10 PARAMETER COMBINATIONS WITH TRADES")
print("="*60)

if 'heatmap' in globals() and heatmap is not None and len(heatmap) > 0:
    # Filter heatmap to include only results where trades actually occurred
    initial_cash = 30000 # Corrected: Use the explicit initial cash amount
    trading_results = heatmap[heatmap != initial_cash].sort_values(ascending=False)

    if not trading_results.empty:
        top_10 = trading_results.head(10)

        print("\nBest 10 combinations by Equity Final [$] (that made trades):\n")

        for i, (idx, value) in enumerate(top_10.items(), 1):
            print(f"#{i} | Equity Final: {value:.2f}")
            print(f"    VIX Low: {idx[0]}, VIX Mid: {idx[1]}, ATR: {idx[2]}")
            print(f"    Risk: {idx[3]*100:.1f}%, Stop: {idx[4]}x, Target: {idx[5]*100:.0f}%")
            print()

        # Also show worst 5 trading results (changed from 3 to 5)
        print("="*60)
        print("WORST 5 TRADING COMBINATIONS (for comparison)")
        print("="*60 + "\n")

        worst_5 = trading_results.sort_values(ascending=True).head(5) # Changed to head(5)

        for i, (idx, value) in enumerate(worst_5.items(), 1):
            print(f"#{i} | Equity Final: {value:.2f}")
            print(f"    VIX Low: {idx[0]}, VIX Mid: {idx[1]}, ATR: {idx[2]}")
            print(f"    Risk: {idx[3]*100:.1f}%, Stop: {idx[4]}x, Target: {idx[5]*100:.0f}%")
            print()
    else:
        print("No parameter combinations resulted in any trades within the explored range.")

else:
    print("\nNo heatmap data available, likely due to optimization failure or `return_heatmap` being False.")

# @title
# Heatmaps
from backtesting.lib import plot_heatmaps
plot_heatmaps(heatmap, agg='mean')

# @title
# Plot
from sambo.plot import plot_evaluations

# Define the names of the parameters for plotting
names = [
    "VIX Low Threshold",
    "VIX Mid Threshold",
    "ATR Period",
    "Risk Per Trade",
    "Stop Multiplier",
    "First Target Portion"
]
_ = plot_evaluations(optimize_result, names=names)

# @title
from sambo.plot import plot_objective
names = ["VIX Low Threshold",
    "VIX Mid Threshold",
    "ATR Period",
    "Risk Per Trade",
    "Stop Multiplier",
    "First Target Portion"]
_ = plot_objective(optimize_result, names=names, estimator='et')

# @title
# Mean Reversion Futures Strategy - CORRECTED VERSION
# Fixed profit target exits and position management
# Compatible with backtesting.py framework

import pandas as pd
import numpy as np
from backtesting import Strategy
from backtesting.lib import crossover


# ========================
# INDICATOR FUNCTIONS
# ========================

def calculate_vwap(close, volume, high, low, index):
    """Calculate VWAP anchored to session start"""
    df = pd.DataFrame({
        'close': close,
        'volume': volume,
        'high': high,
        'low': low
    }, index=index)

    df['typical_price'] = (df['high'] + df['low'] + df['close']) / 3
    df['tp_volume'] = df['typical_price'] * df['volume']

    # Reset at session start (assumes daily sessions)
    df['date'] = df.index.date
    df['cum_tp_volume'] = df.groupby('date')['tp_volume'].cumsum()
    df['cum_volume'] = df.groupby('date')['volume'].cumsum()

    vwap = df['cum_tp_volume'] / df['cum_volume']
    return vwap.values


def calculate_vwap_bands(close, volume, high, low, index, std_mult):
    """Calculate VWAP standard deviation bands"""
    df = pd.DataFrame({
        'close': close,
        'volume': volume,
        'high': high,
        'low': low
    }, index=index)

    df['typical_price'] = (df['high'] + df['low'] + df['close']) / 3
    df['tp_volume'] = df['typical_price'] * df['volume']

    df['date'] = df.index.date
    df['cum_tp_volume'] = df.groupby('date')['tp_volume'].cumsum()
    df['cum_volume'] = df.groupby('date')['volume'].cumsum()

    vwap = df['cum_tp_volume'] / df['cum_volume']

    # Calculate standard deviation
    df['vwap'] = vwap
    df['squared_diff'] = (df['typical_price'] - df['vwap']) ** 2
    df['weighted_sq_diff'] = df['squared_diff'] * df['volume']
    df['cum_weighted_sq_diff'] = df.groupby('date')['weighted_sq_diff'].cumsum()

    variance = df['cum_weighted_sq_diff'] / df['cum_volume']
    std = np.sqrt(variance)

    upper_band = vwap + (std * std_mult)
    lower_band = vwap - (std * std_mult)

    return upper_band.values, lower_band.values


def calculate_atr(high, low, close, period):
    """Calculate Average True Range"""
    df = pd.DataFrame({'high': high, 'low': low, 'close': close})

    df['h_l'] = df['high'] - df['low']
    df['h_pc'] = abs(df['high'] - df['close'].shift(1))
    df['l_pc'] = abs(df['low'] - df['close'].shift(1))

    df['tr'] = df[['h_l', 'h_pc', 'l_pc']].max(axis=1)
    atr = df['tr'].rolling(window=period).mean()

    return atr.values


# ========================
# STRATEGY CLASS
# ========================

class MeanReversionStrategy(Strategy):
    """
    Mean Reversion Futures Strategy - CORRECTED

    Entry Logic:
    - VIX regime filtering (low/mid volatility)
    - Price touches VWAP bands (1σ or 2σ)

    Exit Logic:
    - Profit Target: $100 USD per contract
    - Optional: Stop Loss (3x ATR)

    Position Management:
    - 1 position at a time (backtesting.py limitation)
    """

    # ========================
    # STRATEGY PARAMETERS
    # ========================

    # VIX Regime Thresholds
    vix_low_threshold = 12      # Below this: use 1σ and 2σ bands
    vix_mid_threshold = 27      # Above this: no entries

    # VWAP Bands
    bb_std_1 = 1.0             # First standard deviation
    bb_std_2 = 2.0             # Second standard deviation

    # ATR Settings
    atr_period = 14            # ATR calculation period
    stop_multiplier = 2.0      # Stop loss distance (ATR multiplier)

    # Profit Target
    profit_target_usd = 100.0  # Take profit at $100 per contract

    # Contract Specifications
    price_per_point = 2.0      # MNQ=2, MES=5, ES=50

    # Exit Rules
    use_stop_loss = True      # Enable/disable stop loss
    use_time_exit = False       # Exit after max holding period
    max_holding_bars = 48      # Maximum bars to hold (48h on 1H chart)


    def init(self):
        """Initialize indicators and tracking variables"""

        # Validate VIX data
        if 'VIX' not in self.data.df.columns:
            raise ValueError("ERROR: VIX data not found in DataFrame. "
                           "Ensure your data includes a 'VIX' column.")

        print("=" * 60)
        print("INITIALIZING MEAN REVERSION STRATEGY")
        print("=" * 60)

        # VWAP
        self.vwap = self.I(
            calculate_vwap,
            self.data.Close, self.data.Volume,
            self.data.High, self.data.Low,
            self.data.index,
            name='VWAP'
        )

        # VWAP Bands - 1σ
        bb_result_1 = self.I(
            calculate_vwap_bands,
            self.data.Close, self.data.Volume,
            self.data.High, self.data.Low,
            self.data.index,
            self.bb_std_1,
            name='BB1'
        )
        self.bb_upper_1 = bb_result_1[0]
        self.bb_lower_1 = bb_result_1[1]

        # VWAP Bands - 2σ
        bb_result_2 = self.I(
            calculate_vwap_bands,
            self.data.Close, self.data.Volume,
            self.data.High, self.data.Low,
            self.data.index,
            self.bb_std_2,
            name='BB2'
        )
        self.bb_upper_2 = bb_result_2[0]
        self.bb_lower_2 = bb_result_2[1]

        # ATR
        self.atr = self.I(
            calculate_atr,
            self.data.High, self.data.Low, self.data.Close,
            self.atr_period,
            name='ATR'
        )

        # VIX data
        self.vix = self.data.df['VIX'].values

        # Tracking
        self.current_bar = 0
        self.entry_bar = None
        self.entry_price = None
        self.profit_target_price = None
        self.stop_loss_price = None

        print(f"✓ Contract: {self.price_per_point} USD per point")
        print(f"✓ Profit Target: ${self.profit_target_usd}")
        print(f"✓ Stop Loss: {'Enabled' if self.use_stop_loss else 'DISABLED'}")
        print(f"✓ Time Exit: {self.max_holding_bars} bars" if self.use_time_exit else "✓ Time Exit: Disabled")
        print("=" * 60)


    def next(self):
        """Main strategy logic executed on each bar"""

        self.current_bar += 1

        # Skip if indicators not ready
        if np.isnan(self.vwap[-1]) or np.isnan(self.atr[-1]):
            return

        current_price = self.data.Close[-1]

        # ===========================
        # EXIT MANAGEMENT (IF IN POSITION)
        # ===========================

        if self.position:
            self._check_exits(current_price)
            return  # Don't enter new position if already in one

        # ===========================
        # ENTRY LOGIC (IF NO POSITION)
        # ===========================

        vix_current = self.vix[-1]
        atr_current = self.atr[-1]

        # VIX filter: No entries if VIX too high
        if vix_current > self.vix_mid_threshold:
            return

        # Determine which bands to use based on VIX
        if vix_current < self.vix_low_threshold:
            # Low VIX: Use both 1σ and 2σ bands
            entry_band_1 = self.bb_lower_1[-1]
            entry_band_2 = self.bb_lower_2[-1]
        else:
            # Mid VIX: Only use 2σ band
            entry_band_1 = None
            entry_band_2 = self.bb_lower_2[-1]

        # Check if price touched entry bands
        price_at_band_1 = (entry_band_1 is not None and
                          current_price <= entry_band_1)
        price_at_band_2 = current_price <= entry_band_2

        if not (price_at_band_1 or price_at_band_2):
            return

        # ===========================
        # EXECUTE ENTRY
        # ===========================

        self.buy(size=1)

        # Store entry data for exit management
        self.entry_bar = self.current_bar
        self.entry_price = current_price

        # Calculate profit target price
        profit_target_points = self.profit_target_usd / self.price_per_point
        self.profit_target_price = current_price + profit_target_points

        # Calculate stop loss price (if enabled)
        if self.use_stop_loss:
            stop_distance_points = atr_current * self.stop_multiplier
            self.stop_loss_price = current_price - stop_distance_points
        else:
            self.stop_loss_price = None


    def _check_exits(self, current_price):
        """
        Check all exit conditions for open position

        Exit Priority:
        1. Stop Loss (if enabled)
        2. Profit Target
        3. Time Exit (if enabled)
        """


        # ===========================
        # EXIT 1: STOP LOSS
        # ===========================
        if self.use_stop_loss and self.stop_loss_price is not None:
            if current_price <= self.stop_loss_price:
                self.position.close()
                return

        # ===========================
        # EXIT 2: PROFIT TARGET
        # ===========================
        if self.profit_target_price is not None:
            if current_price >= self.profit_target_price:
                self.position.close()
                return

        # ===========================
        # EXIT 3: TIME-BASED EXIT
        # ===========================
        if self.use_time_exit and self.entry_bar is not None:
            bars_in_trade = self.current_bar - self.entry_bar

            if bars_in_trade >= self.max_holding_bars:
                # Exit if held too long and not profitable enough
                unrealized_pnl_usd = (current_price - self.entry_price) * self.price_per_point

                # Exit if underwater or small profit after long hold
                if unrealized_pnl_usd < self.profit_target_usd * 0.5:  # Less than 50% of target
                    self.position.close()
                    return

# @title
# Backtest
import pandas as pd # Ensure pandas is imported for set_option

# Define custom commission function for a flat fee per order
def custom_commission(size, price):
    # $0.25 flat fee per order (buy or sell)
    return 0.25

bt = Backtest(
    data,
    MeanReversionStrategy,
    cash=30000,
    commission=custom_commission, # Use the custom callable for commission
    margin=0.05,
    trade_on_close=False,
    hedging=False,
    exclusive_orders=True
)

# Run backtest
stats = bt.run()

#Results table
print("\n" + "="*60)
print("BACKTEST RESULTS")
print("="*60)
print(stats)

# Trades
# Export trade log
trades = stats['_trades']

# Columns to drop
columns_to_drop = [
    'Entry_Hammer', 'Exit_Hammer',
    'Entry_Engulfing', 'Exit_Engulfing',
    'Entry_Doji', 'Exit_Doji'
]

# Drop columns if they exist
trades = trades.drop(columns=[col for col in columns_to_drop if col in trades.columns])

# Calculate PnL in USD
# Get the price_per_point directly from the strategy class
price_per_point = MeanReversionStrategy.price_per_point
trades['PnL_USD'] = trades['PnL'] * price_per_point

print("\n" + "="*60)
print(f"Total trades executed: {len(trades)}")
print("="*60)
print("Trade log (all trades):")

# Set display options to show all columns and prevent wrapping
pd.set_option('display.max_columns', None)
pd.set_option('display.width', 1000) # Adjust as needed for your screen
pd.set_option('display.max_rows', None) # Show all rows

print(trades)

# Reset display options to default after printing (optional, but good practice)
pd.reset_option('display.max_columns')
pd.reset_option('display.width')
pd.reset_option('display.max_rows')

# Save to CSV
trades.to_csv('trades.csv')
print("\n Trades exported to trades.csv")