import pandas as pd
import numpy as np
import itertools
import os
import sys
import re
import random
import importlib.util
import argparse
import json
import time
import datetime
import subprocess
import heapq
import sqlite3

# Add project root to sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from strategies import library_activation_scores as lib
from config import RESULTS_DIR, WINNERS_DIR, SWEEPS_DIR, REPORTS_DIR, SCORE_START, TRAIN_END, MVRV_BEAR_THRESHOLD, MVRV_BULL_THRESHOLD, composite_score as _composite_score, get_min_trades as _get_min_trades

# ---------------------------------------------------------------------------
# Optuna-informed biased search configuration
# ---------------------------------------------------------------------------
# Optuna ask/tell is ~1300 trials/s in Python — far too slow to drive our
# 640K samples/s GPU. Instead we use Optuna for two things:
#   1. Post-run importance analysis (offline, fast — only ~5K trials)
#   2. Biased sampling: extract good param regions from DB history, then
#      sample 50% from those regions + 50% pure random (numpy speed, no overhead)
# The study itself is built from sweep DB results and written to a JSON file
# for persistence and the importance report.

SWEEP_DB_FILE      = "results/sweep_database.db"
OPTUNA_STUDIES_DIR = "results/optuna_studies"
DB_SCHEMA_VER      = 3          # must match auto_optimize_loop.py
EXPLOIT_FRACTION   = 0.50       # fraction of samples biased toward good regions
OPTUNA_TOP_K       = 5_000      # top DB rows used for biased-range computation
OPTUNA_TOP_FRAC    = 0.10       # top fraction of those rows used for range bounds
OPTUNA_REFRESH     = 5_000_000  # refresh biased ranges every N GPU samples
GPU_BATCH_SIZE     = 4_000_000  # GPU mini-batch size (1M default; 4M fits ~50% VRAM)

METRIC_TO_DB_COL = {
    "Calmar Ratio":   "calmar_ratio",
    "Sortino Ratio":  "sortino_ratio",
    "Sharpe Ratio":   "sharpe_ratio",
    "P&L/DD Ratio":   "pnl_dd_ratio",
    "Composite":      "composite_score",
}


def _apply_regime_mask(df, regime, bear_threshold, bull_threshold):
    """Filter df to only bars matching the requested regime using the 'zscore' column.
    Returns filtered copy. If 'zscore' column is absent, warns and returns df unchanged."""
    if regime == "all":
        return df
    if "zscore" not in df.columns:
        print(f"[WARNING] --regime={regime} requested but 'zscore' column not found in data; no regime filter applied.")
        return df
    z = df["zscore"]
    if regime == "bear":
        mask = z < bear_threshold
    elif regime == "bull":
        mask = z >= bull_threshold
    else:  # sideways
        mask = (z >= bear_threshold) & (z < bull_threshold)
    n_before = len(df)
    df = df.loc[mask].copy().reset_index(drop=True)
    print(f"[regime={regime}] bear_thr={bear_threshold} bull_thr={bull_threshold}: "
          f"kept {len(df)}/{n_before} IS bars ({len(df)/n_before*100:.1f}%)")
    return df


def _load_db_top_results(asset, timeframe, db_col, top_k=OPTUNA_TOP_K):
    """Return a DataFrame of top_k rows for this asset/TF from the sweep DB."""
    if not os.path.exists(SWEEP_DB_FILE):
        return pd.DataFrame()
    try:
        conn = sqlite3.connect(SWEEP_DB_FILE)
        df = pd.read_sql(
            f"SELECT * FROM sweep_results "
            f"WHERE asset=? AND timeframe=? AND schema_ver=? AND {db_col}>0 "
            f"ORDER BY {db_col} DESC LIMIT ?",
            conn, params=(asset, timeframe, DB_SCHEMA_VER, top_k)
        )
        conn.close()
        return df
    except Exception:
        return pd.DataFrame()


def _compute_biased_ranges(df_top, param_names, param_vals_arrays, top_frac=OPTUNA_TOP_FRAC):
    """
    Compute biased sampling ranges from the top fraction of historical results.
    Returns dict: param_name -> (biased_vals_array | None).
    biased_vals_array contains the subset of the param grid that falls in the p5-p95
    range of the top performers.  None means no meaningful signal (use full range).
    """
    if df_top.empty or len(df_top) < 20:
        return {}

    n_top = max(10, int(len(df_top) * top_frac))
    top_subset = df_top.head(n_top)

    biased = {}
    for name, vals in zip(param_names, param_vals_arrays):
        col = top_subset.get(name)
        if col is None or col.isna().all():
            continue
        col_clean = col.dropna().values.astype(float)
        if len(col_clean) < 5:
            continue
        p5, p95 = np.percentile(col_clean, [5, 95])
        biased_vals = vals[(vals >= p5) & (vals <= p95)]
        if len(biased_vals) < 2:
            continue
        biased[name] = biased_vals
    return biased


def _build_optuna_distributions(param_names, param_vals_arrays, params_config):
    """Build Optuna distribution objects for all non-locked params."""
    try:
        from optuna.distributions import FloatDistribution, CategoricalDistribution
    except ImportError:
        return {}

    dists = {}
    for name, config in params_config.items():
        if "values" in config:
            vals = config["values"]
            if len(vals) > 1:
                dists[name] = CategoricalDistribution([float(v) for v in vals])
        else:
            vals = np.arange(config["start"], config["stop"], config["step"])
            if len(vals) > 0:
                dists[name] = FloatDistribution(
                    float(config["start"]),
                    float(vals[-1]),
                    step=float(config["step"])
                )
    return dists


def _build_optuna_study_from_results(df_top, param_names, params_config, db_col, locked_params):
    """
    Create an in-memory Optuna study pre-populated with top DB results.
    Used only for importance analysis at the end of a run.
    Returns the study, or None if optuna is unavailable / not enough data.
    """
    try:
        import optuna
        optuna.logging.set_verbosity(optuna.logging.WARNING)
    except ImportError:
        return None

    if df_top.empty or len(df_top) < 30:
        return None

    dists = _build_optuna_distributions(param_names, None, params_config)
    # Remove locked params from distributions
    for name in locked_params:
        dists.pop(name, None)
    if not dists:
        return None

    study = optuna.create_study(direction="maximize")
    trials = []
    for _, row in df_top.iterrows():
        score = row.get(db_col)
        if score is None or (isinstance(score, float) and np.isnan(score)) or score <= 0:
            continue
        params = {}
        valid = True
        for name, dist in dists.items():
            val = row.get(name)
            if val is None or (isinstance(val, float) and np.isnan(val)):
                valid = False
                break
            val = float(val)
            # Snap float values onto the grid
            from optuna.distributions import FloatDistribution
            if isinstance(dist, FloatDistribution) and dist.step is not None:
                n_steps = round((val - dist.low) / dist.step)
                val = round(dist.low + n_steps * dist.step, 8)
                val = max(dist.low, min(dist.high, val))
            params[name] = val
        if not valid:
            continue
        try:
            trials.append(optuna.trial.create_trial(
                params=params, distributions=dists, value=float(score)
            ))
        except Exception:
            continue

    if len(trials) < 30:
        return None
    study.add_trials(trials)
    return study


def _write_importance_report(study, asset, timeframe, metric_name):
    """Write param importance report to results/reports/."""
    try:
        import optuna
        # MDI evaluator is ~30× faster than fANOVA and gives comparable rankings
        evaluator = optuna.importance.MeanDecreaseImpurityImportanceEvaluator(seed=42)
        importances = optuna.importance.get_param_importances(study, evaluator=evaluator)
    except Exception as e:
        print(f"   [Optuna] Importance analysis failed: {e}")
        return

    os.makedirs(REPORTS_DIR, exist_ok=True)
    out_path = os.path.join(REPORTS_DIR, f"optuna_importance_{asset}_{timeframe}.md")

    lines = [
        f"# Optuna Parameter Importance — {asset} {timeframe}",
        f"*Generated: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M')}  "
        f"| Metric: {metric_name}  | Trials: {len(study.trials)}*\n",
        "| Rank | Parameter | Importance |",
        "|------|-----------|------------|",
    ]
    for rank, (name, imp) in enumerate(importances.items(), 1):
        bar = "█" * max(1, round(imp * 30))
        lines.append(f"| {rank} | `{name}` | {imp:.4f} {bar} |")

    with open(out_path, "w") as f:
        f.write("\n".join(lines) + "\n")
    print(f"   [Optuna] Importance report → {out_path}")

try:
    import psutil
    from numba import cuda
except ImportError:
    psutil = None
    cuda = None

def get_hardware_stats():
    """Retrieves CPU, RAM, and GPU stats."""
    stats = []
    if psutil:
        stats.append(f"CPU: {psutil.cpu_percent(interval=None):3.0f}%")
        stats.append(f"RAM: {psutil.virtual_memory().percent:3.0f}%")
    try:
        cmd = ['nvidia-smi', '--query-gpu=temperature.gpu,utilization.gpu,memory.used,memory.total', '--format=csv,noheader,nounits']
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=0.1)
        if proc.returncode == 0:
            temp, util, mem_used, mem_total = [int(x) for x in proc.stdout.strip().split(',')]
            stats.append(f"GPU: {temp}C {util}% | VRAM: {(mem_used/mem_total)*100:.0f}%")
    except: pass
    return " | ".join(stats)

def format_large_number(n):
    """Formats large numbers into human-readable strings."""
    if n < 1_000_000: return f"{n:,}"
    if n >= 1e12: return f"{n/1e12:.2f} trillion"
    if n >= 1e9: return f"{n/1e9:.2f} billion"
    return f"{n/1e6:.0f}m"

PROGRESS_INTERVAL = 30.0  # seconds between progress line updates
_last_progress_time = 0.0

def print_progress(iteration, total, start_time, best_score, baseline_score):
    global _last_progress_time
    now = time.time()
    is_final = iteration >= total
    if not is_final and (now - _last_progress_time) < PROGRESS_INTERVAL:
        return
    _last_progress_time = now
    elapsed = now - start_time
    rate = iteration / elapsed if elapsed > 0 else 0
    remaining = (total - iteration) / rate if rate > 0 else 0
    eta = str(datetime.timedelta(seconds=int(remaining)))
    rate_str = f"{rate:,.0f} it/s"
    line = f"Progress: {iteration/total*100:5.1f}% ({format_large_number(iteration)}/{format_large_number(total)}) | Rate: {rate_str} | ETA: {eta} | Best GPU Calmar: {best_score:.4f} (vs {baseline_score:.4f}) | {get_hardware_stats()}"
    sys.stdout.write(f"\r{line}    ")
    if is_final:
        sys.stdout.write("\n")
    sys.stdout.flush()

def main():
    parser = argparse.ArgumentParser(description="Optimize Strategy Parameters")
    parser.add_argument("--data", type=str, required=True, help="Path to data CSV")
    parser.add_argument("--strategy_file", type=str, required=True, help="Name of the strategy Python file")
    parser.add_argument("--params_file", type=str, help="Path to JSON file defining parameter ranges")
    parser.add_argument("--random_search", type=int, default=10000, help="Number of random combinations to try")
    parser.add_argument("--gpu", action="store_true", help="Enable GPU acceleration")
    parser.add_argument("--metric", type=str, default="Calmar Ratio", choices=['Sharpe Ratio', 'Sortino Ratio', 'Calmar Ratio', 'P&L/DD Ratio', 'Composite'], help="The metric to optimize for. 'Composite' = Calmar*(1+log_bonus); GPU pre-filters by raw Calmar, composite applied at CPU verification step.")
    parser.add_argument("--baseline-score", type=float, default=0.0, help="The baseline score for the chosen metric to beat.")
    parser.add_argument("--search", type=str, default="random", choices=["random", "optuna"],
                        help="Search strategy: 'random' (pure uniform) or 'optuna' (biased sampling "
                             "from DB history + post-run importance analysis).")
    parser.add_argument("--timeframe", type=str, default=None, help="Timeframe suffix for output files (e.g. '1D'). Auto-derived from --data filename if not set.")
    parser.add_argument("--asset", type=str, default=None, help="Asset name for output files (e.g. 'COINBASE_BTCUSD'). Auto-derived from --data filename if not set.")
    parser.add_argument("--regime", type=str, default="all", choices=["all", "bull", "bear", "sideways"],
                        help="Pre-filter IS training bars by MVRV regime. "
                             "'bull': zscore >= bull_threshold; 'bear': zscore < bear_threshold; "
                             "'sideways': bear_threshold <= zscore < bull_threshold; 'all': no filter (default).")
    parser.add_argument("--regime-bear-threshold", type=float, default=MVRV_BEAR_THRESHOLD,
                        help=f"zscore percentile below which bars are classified as Bear (default: {MVRV_BEAR_THRESHOLD}).")
    parser.add_argument("--regime-bull-threshold", type=float, default=MVRV_BULL_THRESHOLD,
                        help=f"zscore percentile at or above which bars are classified as Bull (default: {MVRV_BULL_THRESHOLD}).")
    args = parser.parse_args()

    # Derive asset and timeframe from data filename.
    # Supports two naming formats:
    #   mlp format:  "COINBASE_BTCUSD, 240.csv"  -> asset=COINBASE_BTCUSD, tf=4H
    #   legacy format: "COINBASE_BTCUSD-1D.csv"  -> asset=COINBASE_BTCUSD, tf=1D
    _MIN_TO_TF = {"240": "4H", "360": "6H", "480": "8H", "720": "12H", "1D": "1D"}
    stem = os.path.splitext(os.path.basename(args.data))[0]
    if ', ' in stem:
        asset_part, period = stem.rsplit(', ', 1)
        derived_asset, derived_tf = asset_part, _MIN_TO_TF.get(period, period.upper())
    else:
        parts = stem.rsplit('-', 1)
        derived_asset = parts[0] if len(parts) == 2 else stem
        derived_tf = parts[1].upper() if len(parts) == 2 else "1D"
    if not args.asset:
        args.asset = derived_asset
    if not args.timeframe:
        args.timeframe = derived_tf
    args.asset_tf = f"{args.asset}_{args.timeframe}"

    # --- Dynamic Strategy Import ---
    script_dir = os.path.dirname(os.path.abspath(__file__))
    strategy_path = os.path.join(script_dir, os.path.basename(args.strategy_file))

    if not os.path.exists(strategy_path):
        strategy_path = args.strategy_file
        if not os.path.exists(strategy_path):
             print(f"Error: Strategy file not found at {strategy_path}")
             return

    spec = importlib.util.spec_from_file_location("strategy_module", strategy_path)
    strategy_signals = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(strategy_signals)
    
    print(f"Loading data from {args.data}...", flush=True)
    df = pd.read_csv(args.data)
    df.columns = df.columns.str.lower().str.strip()
    
    # --- Parameter Grid Setup ---
    strategy_basename = os.path.splitext(os.path.basename(args.strategy_file))[0]
    strategy_name = strategy_basename.replace("strategy_", "")
    params_file = args.params_file or f"strategies/params_{strategy_basename}.json"
    with open(params_file, 'r') as f:
        params_config = json.load(f)

    param_names, param_values_list = [], []
    for name, config in params_config.items():
        param_names.append(name)
        if "values" in config:
            param_values_list.append(config["values"])
        elif "start" in config:
            vals = np.arange(config["start"], config["stop"], config["step"])
            param_values_list.append([round(x, 8) for x in vals] or [config["start"]])

    # When i_regime_window is locked to 0 (regime filter off), i_regime_entry_min_score
    # has no effect on strategy output. Lock it to a constant to avoid wasting search budget.
    if "i_regime_window" in param_names and "i_regime_entry_min_score" in param_names:
        rw_idx = param_names.index("i_regime_window")
        if param_values_list[rw_idx] == [0]:
            rs_idx = param_names.index("i_regime_entry_min_score")
            param_values_list[rs_idx] = [-1000.0]

    # --- Baseline Setup ---
    os.makedirs(RESULTS_DIR, exist_ok=True)
    os.makedirs(WINNERS_DIR, exist_ok=True)
    os.makedirs(SWEEPS_DIR, exist_ok=True)
    winner_filename = os.path.join(WINNERS_DIR, f"optimization_winner_{strategy_name}_{args.asset_tf}.csv")
    initial_best_score = args.baseline_score
    best_score = initial_best_score
    print(f"Attempting to beat baseline Calmar (GPU pre-filter): {initial_best_score:.4f}  [CPU optimising for: {args.metric}]")
    
    # --- Search Execution ---
    total_combos = 1
    for v in param_values_list:
        total_combos *= len(v)
    use_random = args.random_search > 0 and total_combos > args.random_search
    iterations = args.random_search if use_random else total_combos
    use_optuna = args.search == "optuna"
    search_label = ("Optuna-Biased" if use_optuna else ("Random" if use_random else "Grid"))
    print(f"Starting {search_label} Search with {iterations:,} iterations (optimizing for {args.metric})...", flush=True)

    sweep_filename = os.path.join(SWEEPS_DIR, f"optimization_sweep_{strategy_name}_{args.asset_tf}.csv")
    if os.path.exists(sweep_filename): os.remove(sweep_filename)

    best_result = None
    run_best_result = None
    run_best_score = -float('inf')
    start_time = time.time()
    top_results_heap = []
    MAX_SWEEP_RESULTS = 5000

    # --- GPU Path ---
    if args.gpu and hasattr(strategy_signals, 'execute_gpu_batch'):
        BATCH_SIZE = GPU_BATCH_SIZE
        total_processed = 0

        print("Pre-calculating indicators and uploading to GPU...", flush=True)
        df_gpu = df.copy()
        if 'time' in df_gpu.columns:
            df_gpu['time'] = pd.to_datetime(df_gpu['time'], utc=True).dt.tz_localize(None)
            mask = (df_gpu['time'] >= '2015-01-01') & (df_gpu['time'] <= TRAIN_END)
            df_gpu = df_gpu.loc[mask].copy().reset_index(drop=True)
        df_gpu = _apply_regime_mask(df_gpu, args.regime, args.regime_bear_threshold, args.regime_bull_threshold)
        gpu_static_data = strategy_signals.prepare_gpu_data(df_gpu)
        T = len(df_gpu)  # number of daily bars; used for annualising returns (Calmar)

        weight_names = [p for p in param_names if p.startswith('i_w_')]
        threshold_names = [p for p in param_names if 'threshold' in p]
        config_names = [p for p in param_names if p.startswith('i_use_')]
        param_vals_arrays = [np.array(vals) for vals in param_values_list]
        locked_params = {n: v[0] for n, v in zip(param_names, param_values_list) if len(v) == 1}

        # --- Optuna: load DB history and compute initial biased ranges ---
        db_col = METRIC_TO_DB_COL.get(args.metric, "calmar_ratio")
        biased_ranges = {}
        next_refresh_at = OPTUNA_REFRESH
        if use_optuna:
            df_db_top = _load_db_top_results(args.asset, args.timeframe, db_col)
            if not df_db_top.empty:
                biased_ranges = _compute_biased_ranges(df_db_top, param_names, param_vals_arrays)
                n_biased = sum(1 for v in biased_ranges.values() if v is not None)
                print(f"   [Optuna] Loaded {len(df_db_top)} DB rows; biased ranges for {n_biased}/{len(param_names)} params "
                      f"({EXPLOIT_FRACTION*100:.0f}% exploitation).", flush=True)
            else:
                print("   [Optuna] No DB history yet — running as random search this iteration.", flush=True)

        try:
            while total_processed < iterations:
                current_batch_size = min(BATCH_SIZE, iterations - total_processed)

                weights = np.zeros((current_batch_size, len(weight_names)), dtype=np.float64)
                thresholds = np.zeros((current_batch_size, len(threshold_names)), dtype=np.float64)
                configs = np.zeros((current_batch_size, 2), dtype=np.float64)
                batch_values_cols = []

                for idx, (name, vals) in enumerate(zip(param_names, param_vals_arrays)):
                    biased_v = biased_ranges.get(name) if use_optuna else None

                    if biased_v is not None and len(biased_v) >= 2:
                        # Mix: EXPLOIT_FRACTION from biased range, rest from full range
                        n_exploit = int(current_batch_size * EXPLOIT_FRACTION)
                        n_explore = current_batch_size - n_exploit
                        exploit_col = biased_v[np.random.randint(0, len(biased_v), size=n_exploit)]
                        explore_col = vals[np.random.randint(0, len(vals), size=n_explore)]
                        col_values = np.empty(current_batch_size, dtype=np.float64)
                        col_values[:n_exploit] = exploit_col
                        col_values[n_exploit:] = explore_col
                        # Shuffle so exploit/explore rows are intermixed across the batch
                        np.random.shuffle(col_values)
                    else:
                        col_values = vals[np.random.randint(0, len(vals), size=current_batch_size)]

                    batch_values_cols.append(col_values)
                    if name in weight_names:     weights[:, weight_names.index(name)] = col_values
                    elif name in threshold_names: thresholds[:, threshold_names.index(name)] = col_values
                    elif name in config_names:    configs[:, config_names.index(name)] = col_values

                gpu_results = strategy_signals.execute_gpu_batch(gpu_static_data, weights, thresholds, configs)

                for i in range(current_batch_size):
                    pnl, dd, sharpe, sortino, trades = gpu_results[i]

                    # Calmar: annualised return / abs(max drawdown %).
                    # dd from kernel is already abs(max_drawdown*100) — a positive value.
                    if trades >= 30 and dd > 0.1 and (1.0 + pnl / 100.0) > 1e-9:
                        annualised_pct = ((1.0 + pnl / 100.0) ** (365.0 / T) - 1.0) * 100.0
                        calmar = annualised_pct / dd
                    else:
                        calmar = -10.0

                    if args.metric in ('Calmar Ratio', 'Composite'):
                        # NOTE: GPU pre-filter always uses raw Calmar. For 'Composite',
                        # the actual composite score is computed after CPU verification
                        # in auto_optimize_loop.py (requires is_years, unavailable here).
                        current_score = calmar
                    elif args.metric == 'Sortino Ratio':
                        current_score = sortino
                    else:
                        current_score = sharpe
                    if not np.isfinite(current_score): continue

                    if current_score > run_best_score:
                        run_best_score = current_score
                        current_params = {param_names[j]: batch_values_cols[j][i] for j in range(len(param_names))}
                        run_best_result = {**current_params, 'Total P&L %': pnl, 'Max Drawdown %': dd,
                                           'Sharpe Ratio': sharpe, 'Sortino Ratio': sortino,
                                           'Calmar Ratio': calmar, 'Total Trades': trades}

                        if current_score > best_score:
                            best_score = current_score
                            best_result = run_best_result
                            pd.DataFrame([best_result]).to_csv(os.path.join(WINNERS_DIR, f"optimization_run_best_{strategy_name}_{args.asset_tf}.csv"), index=False)

                    if len(top_results_heap) < MAX_SWEEP_RESULTS or current_score > top_results_heap[0][0]:
                        res_dict = {param_names[j]: batch_values_cols[j][i] for j in range(len(param_names))}
                        res_dict.update({'Sharpe Ratio': sharpe, 'Sortino Ratio': sortino,
                                         'Calmar Ratio': calmar, 'Total Trades': trades})
                        if len(top_results_heap) < MAX_SWEEP_RESULTS:
                            heapq.heappush(top_results_heap, (current_score, total_processed + i, res_dict))
                        else:
                            heapq.heappushpop(top_results_heap, (current_score, total_processed + i, res_dict))

                total_processed += current_batch_size

                # --- Optuna: periodically refresh biased ranges from accumulated heap ---
                if use_optuna and total_processed >= next_refresh_at and top_results_heap:
                    next_refresh_at += OPTUNA_REFRESH
                    # Rebuild biased ranges from the current top-heap results
                    heap_records = [item[2] for item in top_results_heap]
                    df_heap = pd.DataFrame(heap_records)
                    df_heap[db_col] = df_heap.get(args.metric, df_heap.get(db_col))
                    new_biased = _compute_biased_ranges(df_heap, param_names, param_vals_arrays)
                    if new_biased:
                        biased_ranges = new_biased
                        n_updated = sum(1 for v in biased_ranges.values() if v is not None)
                        print(f"\n   [Optuna] Refreshed biased ranges at {total_processed:,} samples "
                              f"({n_updated} params directed).", flush=True)

                print_progress(total_processed, iterations, start_time, best_score, initial_best_score)
        except (KeyboardInterrupt, Exception) as e:
            print(f"\nOptimization stopped: {e}")

    # --- CPU Path ---
    else:
        # Pre-filter once before the loop to avoid redundant work per iteration
        df_cpu_base = df.copy()
        if 'time' in df_cpu_base.columns:
            df_cpu_base['time'] = pd.to_datetime(df_cpu_base['time'], utc=True).dt.tz_localize(None)
            mask = (df_cpu_base['time'] >= '2015-01-01') & (df_cpu_base['time'] <= TRAIN_END)
            df_cpu_base = df_cpu_base.loc[mask].copy()
        df_cpu_base = _apply_regime_mask(df_cpu_base, args.regime, args.regime_bear_threshold, args.regime_bull_threshold)

        combinations = (dict(zip(param_names, [random.choice(v) for v in param_values_list])) for _ in range(iterations)) if use_random else (dict(zip(param_names, c)) for c in itertools.product(*param_values_list))
        for i, kwargs in enumerate(combinations):
            df_cpu = df_cpu_base.copy()

            signals = strategy_signals.generate_signals(df_cpu, **kwargs)
            metrics = strategy_signals.calculate_metrics(signals)
            current_score = metrics.get(args.metric, 0.0)
            
            if current_score > run_best_score:
                run_best_score = current_score
                run_best_result = {**kwargs, **metrics}
                if current_score > best_score:
                    best_score = current_score
                    best_result = run_best_result
            
            print_progress(i + 1, iterations, start_time, best_score, initial_best_score)
            if i + 1 >= iterations: break

    # --- Save results ---
    print("\n\nOptimization Complete.")
    if top_results_heap:
        # For 'Composite', the GPU heap is sorted by raw Calmar (composite applied post-CPU-verify).
        gpu_sort_key = 'Calmar Ratio' if args.metric == 'Composite' else args.metric
        sorted_sweep = sorted([item[2] for item in top_results_heap], key=lambda x: x.get(gpu_sort_key, 0), reverse=True)
        pd.DataFrame(sorted_sweep).to_csv(sweep_filename, index=False)
        print(f"Top {len(sorted_sweep)} results saved to {sweep_filename}")

    if run_best_result:
        pd.DataFrame([run_best_result]).to_csv(os.path.join(WINNERS_DIR, f"optimization_run_best_{strategy_name}_{args.asset_tf}.csv"), index=False)
        print(f"Best result of this run saved.")

    if best_result and best_score > initial_best_score:
        # NOTE: do NOT write to winner_filename here. auto_optimize_loop.py is the
        # authoritative writer — it runs CPU verification + WFO after this subprocess
        # exits and always writes the definitive winner. Writing here produces a
        # preliminary Calmar-only result that the parent immediately has to restore over.
        print(f"\n[SUCCESS] GPU found improvement over baseline {args.metric} (pending CPU verification + WFO).")
    else:
        print(f"\n[INFO] No improvement found over baseline {args.metric}.")

    # --- Optuna: post-run importance analysis ---
    if use_optuna and top_results_heap:
        print("\n   [Optuna] Running parameter importance analysis...", flush=True)
        # Combine current run's top results with DB history for richer analysis
        run_records = [item[2] for item in top_results_heap]
        df_run = pd.DataFrame(run_records)
        db_col_local = METRIC_TO_DB_COL.get(args.metric, "calmar_ratio")
        df_run[db_col_local] = df_run.get(args.metric, df_run.get(db_col_local))
        # Cap at 1000 rows for importance analysis — MDI is O(N) but 2K+ starts to feel slow
        df_db_hist = _load_db_top_results(args.asset, args.timeframe, db_col_local, top_k=1000)
        df_combined = pd.concat([df_run.head(500), df_db_hist], ignore_index=True).drop_duplicates()
        locked_params_local = {n: v[0] for n, v in zip(param_names, param_values_list) if len(v) == 1}
        study = _build_optuna_study_from_results(
            df_combined, param_names, params_config, db_col_local, locked_params_local
        )
        if study is not None:
            _write_importance_report(study, args.asset, args.timeframe, args.metric)
        else:
            print("   [Optuna] Not enough data for importance analysis (need ≥30 trials).")

    sys.stdout.flush()
    sys.stderr.flush()
    if cuda:
        cuda.close()
    sys.exit(0)

if __name__ == "__main__":
    main()