"""Pandas vs Polars Benchmark — systematic performance comparison across financial data operations."""

import gc
import warnings

import numpy as np
import pandas as pd
import plotly.graph_objects as go
import polars as pl
import psutil
from IPython.display import display
from plotly.subplots import make_subplots

from utils.reproducibility import set_global_seeds
from utils.storage_benchmarks import (
    ACTIVE_SCALE,
    BENCHMARK_DIR,
    N_ROWS_PER_SYMBOL,
    N_SYMBOLS,
    RESULTS_DIR,
    TIMING_RUNS,
    estimate_memory_mb,
    generate_ohlcv_data,
    generate_tick_data,
    get_scale_config,
    time_operation,
)
from utils.style import COLORS

warnings.filterwarnings("ignore")

# Production defaults — Papermill injects overrides for CI
SEED = 42

set_global_seeds(SEED)

# Version detection
PANDAS_VERSION = pd.__version__
POLARS_VERSION = pl.__version__

print("=" * 70)
print("DATAFRAME LIBRARY BENCHMARK")
print("=" * 70)
print(f"\npandas version: {PANDAS_VERSION}")
print(f"Polars version: {POLARS_VERSION}")

# pandas 3.0 feature detection
PANDAS_MAJOR = int(PANDAS_VERSION.split(".")[0])
IS_PANDAS_3 = PANDAS_MAJOR >= 3

# Check Copy-on-Write status (enabled by default in pandas 3.0)
try:
    COW_ENABLED = pd.options.mode.copy_on_write
except AttributeError:
    COW_ENABLED = False

# Check for PyArrow string dtype (default in pandas 3.0)
try:
    test_series = pd.Series(["test"])
    PYARROW_STRINGS = "pyarrow" in str(test_series.dtype) or test_series.dtype == "string"
except Exception:
    PYARROW_STRINGS = False

print("\npandas 3.0 features:")
print(f"  Copy-on-Write: {'enabled' if COW_ENABLED else 'disabled'}")
print(f"  PyArrow strings: {'yes' if PYARROW_STRINGS else 'no'}")

# Configure Polars streaming (opt-in to new engine in 1.37+)
POLARS_STREAMING = False
try:
    pl.Config.set_engine_affinity(engine="streaming")
    POLARS_STREAMING = True
    print("\nPolars streaming engine: enabled")
except Exception:
    print("\nPolars streaming engine: not available (requires 1.37+)")

scale_cfg = get_scale_config(ACTIVE_SCALE)
print(f"\nScale: {ACTIVE_SCALE} ({scale_cfg['target_memory']} target)")
print(f"OHLCV: {N_SYMBOLS} symbols × {N_ROWS_PER_SYMBOL:,} rows/symbol")

print("\n=== Generating synthetic data ===\n")

# Generate OHLCV data (Polars native)
ohlcv_pl = generate_ohlcv_data(n_symbols=N_SYMBOLS, n_rows=N_ROWS_PER_SYMBOL)
total_rows = len(ohlcv_pl)
print(f"OHLCV: {total_rows:,} rows ({estimate_memory_mb(ohlcv_pl):.1f} MB)")

# Convert to pandas (triggers CoW in pandas 3.0)
ohlcv_pd = ohlcv_pl.to_pandas()
print(f"pandas memory: {ohlcv_pd.memory_usage(deep=True).sum() / 1e6:.1f} MB")

# Generate tick data for join benchmarks
trades_pl, quotes_pl = generate_tick_data(
    n_symbols=min(N_SYMBOLS, 50),  # Limit symbols for tick data
    seed=42,
)
n_trades = len(trades_pl)
n_quotes = len(quotes_pl)
print(f"Trades: {n_trades:,} rows")
print(f"Quotes: {n_quotes:,} rows")

# Convert tick data to pandas
trades_pd = trades_pl.to_pandas()
quotes_pd = quotes_pl.to_pandas()

# Store results
results = []

def force_eval_pandas(df: pd.DataFrame) -> None:
    """Force pandas DataFrame evaluation by touching all data."""
    # Numeric columns: sum
    numeric_cols = df.select_dtypes(include=[np.number]).columns
    if len(numeric_cols) > 0:
        _ = df[numeric_cols].sum().sum()
    # String columns: length (only actual string types)
    str_cols = df.select_dtypes(include=["object", "string"]).columns
    for col in str_cols[:2]:  # Limit to avoid slow string ops
        try:
            if df[col].dtype == "object" or "string" in str(df[col].dtype):
                _ = df[col].astype(str).str.len().sum()
        except Exception:
            pass  # Skip if not actually string-like

def force_eval_polars(df: pl.DataFrame | pl.LazyFrame) -> pl.DataFrame:
    """Force Polars DataFrame evaluation."""
    if isinstance(df, pl.LazyFrame):
        df = df.collect()
    # Touch numeric columns
    numeric_cols = [c for c in df.columns if df[c].dtype in (pl.Float64, pl.Int64)]
    if numeric_cols:
        _ = df.select([pl.col(c).sum() for c in numeric_cols[:5]]).to_dict()
    return df

def benchmark_operation(
    name: str,
    category: str,
    pandas_func,
    polars_func,
    n_runs: int = TIMING_RUNS,
) -> dict:
    """Benchmark an operation on both libraries.

    Returns dict with timing results for both libraries.
    """
    # pandas benchmark
    gc.collect()
    pd_time, pd_result = time_operation(pandas_func, n_runs=n_runs)

    # Polars benchmark
    gc.collect()
    pl_time, pl_result = time_operation(polars_func, n_runs=n_runs)

    # Calculate speedup
    speedup = pd_time / pl_time if pl_time > 0 else float("inf")

    result = {
        "category": category,
        "operation": name,
        "pandas_time": pd_time,
        "polars_time": pl_time,
        "speedup": speedup,
    }

    print(f"  {name}: pandas={pd_time:.4f}s, polars={pl_time:.4f}s, speedup={speedup:.1f}x")

    return result

print("\n" + "=" * 70)
print("CATEGORY A: ROLLING CALCULATIONS")
print("=" * 70)

rolling_results = []

def pd_rolling_mean():
    result = ohlcv_pd.groupby("symbol")["close"].rolling(20).mean().reset_index(drop=True)
    _ = result.sum()  # Force evaluation
    return result

def pl_rolling_mean():
    """Compute 20-day rolling mean per symbol using Polars window expressions."""
    result = ohlcv_pl.with_columns(pl.col("close").rolling_mean(20).over("symbol").alias("sma_20"))
    _ = result.select(pl.col("sma_20").sum()).item()
    return result


r = benchmark_operation("rolling_mean_20", "A_rolling", pd_rolling_mean, pl_rolling_mean)
rolling_results.append(r)

def pd_rolling_std():
    result = ohlcv_pd.groupby("symbol")["close"].rolling(20).std().reset_index(drop=True)
    _ = result.sum()
    return result

def pl_rolling_std():
    """Compute 20-day rolling standard deviation per symbol for volatility estimation."""
    result = ohlcv_pl.with_columns(pl.col("close").rolling_std(20).over("symbol").alias("vol_20"))
    _ = result.select(pl.col("vol_20").sum()).item()
    return result


r = benchmark_operation("rolling_std_20", "A_rolling", pd_rolling_std, pl_rolling_std)
rolling_results.append(r)

HORIZONS = [1, 5, 21, 63, 126, 252]


def pd_multi_horizon():
    result = ohlcv_pd.copy()
    for h in HORIZONS:
        result[f"ret_{h}"] = result.groupby("symbol")["close"].pct_change(h)
    force_eval_pandas(result)
    return result

def pl_multi_horizon():
    """Compute returns at six horizons in a single with_columns call using Polars expressions."""
    # Polars: all horizons in single with_columns call
    result = ohlcv_pl.with_columns(
        [pl.col("close").pct_change(h).over("symbol").alias(f"ret_{h}") for h in HORIZONS]
    )
    _ = result.select([pl.col(f"ret_{h}").sum() for h in HORIZONS]).to_dict()
    return result


r = benchmark_operation("multi_horizon_returns", "A_rolling", pd_multi_horizon, pl_multi_horizon)
rolling_results.append(r)

def pd_rolling_sharpe():
    result = ohlcv_pd.copy()
    returns = result.groupby("symbol")["close"].pct_change()
    result["sharpe"] = (returns.rolling(63).mean() / returns.rolling(63).std()) * np.sqrt(252)
    force_eval_pandas(result)
    return result

def pl_rolling_sharpe():
    """Compute 63-day rolling Sharpe ratio via chained Polars window expressions."""
    result = ohlcv_pl.with_columns(
        pl.col("close").pct_change().over("symbol").alias("returns")
    ).with_columns(
        (
            pl.col("returns").rolling_mean(63).over("symbol")
            / pl.col("returns").rolling_std(63).over("symbol")
        )
        .mul(np.sqrt(252))
        .alias("sharpe")
    )
    _ = result.select(pl.col("sharpe").sum()).item()
    return result


r = benchmark_operation("rolling_sharpe_63", "A_rolling", pd_rolling_sharpe, pl_rolling_sharpe)
rolling_results.append(r)

def pd_ewm():
    result = (
        ohlcv_pd.groupby("symbol")["close"].ewm(span=20, adjust=False).mean().reset_index(drop=True)
    )
    _ = result.sum()
    return result

def pl_ewm():
    """Compute exponential moving average (span=20) per symbol using Polars ewm_mean."""
    result = ohlcv_pl.with_columns(
        pl.col("close").ewm_mean(span=20, adjust=False).over("symbol").alias("ema_20")
    )
    _ = result.select(pl.col("ema_20").sum()).item()
    return result


r = benchmark_operation("ewm_span_20", "A_rolling", pd_ewm, pl_ewm)
rolling_results.append(r)

results.extend(rolling_results)

print("\n" + "=" * 70)
print("CATEGORY B: GROUPBY AGGREGATIONS")
print("=" * 70)

groupby_results = []

def pd_resample():
    result = (
        ohlcv_pd.groupby([ohlcv_pd["timestamp"].dt.date, "symbol"])
        .agg(
            {
                "open": "first",
                "high": "max",
                "low": "min",
                "close": "last",
                "volume": "sum",
            }
        )
        .reset_index()
    )
    force_eval_pandas(result)
    return result

def pl_resample():
    """Resample OHLCV to daily bars using Polars group_by with first/last/min/max/sum aggregations."""
    result = ohlcv_pl.group_by([pl.col("timestamp").dt.date().alias("timestamp"), "symbol"]).agg(
        [
            pl.col("open").first(),
            pl.col("high").max(),
            pl.col("low").min(),
            pl.col("close").last(),
            pl.col("volume").sum(),
        ]
    )
    force_eval_polars(result)
    return result


r = benchmark_operation("ohlcv_resample_daily", "B_groupby", pd_resample, pl_resample)
groupby_results.append(r)

def pd_cross_sectional():
    result = ohlcv_pd.groupby("timestamp").agg(
        {
            "close": ["mean", "std", "min", "max"],
            "volume": ["sum", "mean"],
        }
    )
    result.columns = ["_".join(col) for col in result.columns]
    return result.reset_index()

def pl_cross_sectional():
    """Compute cross-sectional statistics (mean, std, min, max) per timestamp using Polars."""
    result = ohlcv_pl.group_by("timestamp").agg(
        [
            pl.col("close").mean().alias("close_mean"),
            pl.col("close").std().alias("close_std"),
            pl.col("close").min().alias("close_min"),
            pl.col("close").max().alias("close_max"),
            pl.col("volume").sum().alias("volume_sum"),
            pl.col("volume").mean().alias("volume_mean"),
        ]
    )
    force_eval_polars(result)
    return result


r = benchmark_operation(
    "cross_sectional_stats", "B_groupby", pd_cross_sectional, pl_cross_sectional
)
groupby_results.append(r)

def pd_symbol_stats():
    result = ohlcv_pd.groupby("symbol").agg(
        {
            "close": ["mean", "std", "min", "max", "count"],
            "volume": ["sum", "mean"],
            "high": "max",
            "low": "min",
        }
    )
    result.columns = ["_".join(col) for col in result.columns]
    return result.reset_index()

def pl_symbol_stats():
    """Compute per-symbol summary statistics (close, volume, high, low) using Polars group_by."""
    result = ohlcv_pl.group_by("symbol").agg(
        [
            pl.col("close").mean().alias("close_mean"),
            pl.col("close").std().alias("close_std"),
            pl.col("close").min().alias("close_min"),
            pl.col("close").max().alias("close_max"),
            pl.col("close").count().alias("close_count"),
            pl.col("volume").sum().alias("volume_sum"),
            pl.col("volume").mean().alias("volume_mean"),
            pl.col("high").max().alias("high_max"),
            pl.col("low").min().alias("low_min"),
        ]
    )
    force_eval_polars(result)
    return result


r = benchmark_operation("symbol_stats", "B_groupby", pd_symbol_stats, pl_symbol_stats)
groupby_results.append(r)

results.extend(groupby_results)

print("\n" + "=" * 70)
print("CATEGORY C: WINDOW FUNCTIONS")
print("=" * 70)

window_results = []

def pd_zscore():
    result = ohlcv_pd.copy()
    result["returns"] = result.groupby("symbol")["close"].pct_change()
    grouped = result.groupby("timestamp")["returns"]
    result["zscore"] = (result["returns"] - grouped.transform("mean")) / grouped.transform("std")
    force_eval_pandas(result)
    return result

def pl_zscore():
    """Compute cross-sectional z-score of returns at each timestamp using Polars .over() windows."""
    result = ohlcv_pl.with_columns(
        pl.col("close").pct_change().over("symbol").alias("returns")
    ).with_columns(
        (
            (pl.col("returns") - pl.col("returns").mean().over("timestamp"))
            / pl.col("returns").std().over("timestamp")
        ).alias("zscore")
    )
    _ = result.select(pl.col("zscore").sum()).item()
    return result


r = benchmark_operation("cross_sectional_zscore", "C_window", pd_zscore, pl_zscore)
window_results.append(r)

def pd_rank():
    result = ohlcv_pd.copy()
    result["returns"] = result.groupby("symbol")["close"].pct_change()
    result["rank"] = result.groupby("timestamp")["returns"].rank(pct=True)
    force_eval_pandas(result)
    return result

def pl_rank():
    """Compute percentile rank of returns within each timestamp cross-section using Polars."""
    result = (
        ohlcv_pl.with_columns(pl.col("close").pct_change().over("symbol").alias("returns"))
        .with_columns(pl.col("returns").rank().over("timestamp").alias("rank_raw"))
        .with_columns(
            (pl.col("rank_raw") / pl.col("rank_raw").max().over("timestamp")).alias("rank_pct")
        )
    )
    _ = result.select(pl.col("rank_pct").sum()).item()
    return result


r = benchmark_operation("percentile_rank", "C_window", pd_rank, pl_rank)
window_results.append(r)

LAGS = [1, 5, 21]


def pd_lags():
    result = ohlcv_pd.copy()
    for lag in LAGS:
        result[f"close_lag_{lag}"] = result.groupby("symbol")["close"].shift(lag)
    force_eval_pandas(result)
    return result

def pl_lags():
    """Create multiple lag columns (1, 5, 21 days) per symbol using Polars shift with .over()."""
    result = ohlcv_pl.with_columns(
        [pl.col("close").shift(lag).over("symbol").alias(f"close_lag_{lag}") for lag in LAGS]
    )
    _ = result.select([pl.col(f"close_lag_{lag}").sum() for lag in LAGS]).to_dict()
    return result


r = benchmark_operation("lagged_values", "C_window", pd_lags, pl_lags)
window_results.append(r)

results.extend(window_results)

print("\n" + "=" * 70)
print("CATEGORY D: FILTERING")
print("=" * 70)

filter_results = []

# Precompute price threshold (median)
price_threshold = float(ohlcv_pl.select(pl.col("close").median()).item())


def pd_simple_filter():
    result = ohlcv_pd[ohlcv_pd["close"] > price_threshold]
    force_eval_pandas(result)
    return result

def pl_simple_filter():
    """Filter rows where close exceeds median price threshold using Polars filter."""
    result = ohlcv_pl.filter(pl.col("close") > price_threshold)
    force_eval_polars(result)
    return result


r = benchmark_operation("simple_filter", "D_filter", pd_simple_filter, pl_simple_filter)
filter_results.append(r)

# Get list of symbols for filtering
symbol_list = ohlcv_pl.select("symbol").unique().head(N_SYMBOLS // 2)["symbol"].to_list()
volume_threshold = float(ohlcv_pl.select(pl.col("volume").median()).item())


def pd_multi_filter():
    result = ohlcv_pd[
        (ohlcv_pd["close"] > price_threshold)
        & (ohlcv_pd["volume"] > volume_threshold)
        & (ohlcv_pd["symbol"].isin(symbol_list))
    ]
    force_eval_pandas(result)
    return result

def pl_multi_filter():
    """Apply compound filter on price, volume, and symbol membership using Polars boolean expressions."""
    result = ohlcv_pl.filter(
        (pl.col("close") > price_threshold)
        & (pl.col("volume") > volume_threshold)
        & (pl.col("symbol").is_in(symbol_list))
    )
    force_eval_polars(result)
    return result


r = benchmark_operation("multi_condition_filter", "D_filter", pd_multi_filter, pl_multi_filter)
filter_results.append(r)

def pd_range_filter():
    # Simulate: price between 95-105% of reference, volume in range
    ref_price = price_threshold
    result = ohlcv_pd[
        (ohlcv_pd["close"] >= ref_price * 0.95)
        & (ohlcv_pd["close"] <= ref_price * 1.05)
        & (ohlcv_pd["volume"] >= volume_threshold * 0.5)
        & (ohlcv_pd["volume"] <= volume_threshold * 2.0)
    ]
    force_eval_pandas(result)
    return result

def pl_range_filter():
    """Filter by price and volume ranges using Polars is_between for options-style moneyness bands."""
    ref_price = price_threshold
    result = ohlcv_pl.filter(
        pl.col("close").is_between(ref_price * 0.95, ref_price * 1.05)
        & pl.col("volume").is_between(volume_threshold * 0.5, volume_threshold * 2.0)
    )
    force_eval_polars(result)
    return result


r = benchmark_operation("range_filter", "D_filter", pd_range_filter, pl_range_filter)
filter_results.append(r)

results.extend(filter_results)

print("\n" + "=" * 70)
print("CATEGORY E: JOINS")
print("=" * 70)

join_results = []

# Sort data for ASOF join (required)
# pandas merge_asof requires left keys to be sorted by the "on" column
trades_pd_sorted = trades_pd.sort_values("timestamp").reset_index(drop=True)
quotes_pd_sorted = quotes_pd.sort_values("timestamp").reset_index(drop=True)
# Polars requires sort by both by and on columns
trades_pl_sorted = trades_pl.sort(["symbol", "timestamp"])
quotes_pl_sorted = quotes_pl.sort(["symbol", "timestamp"])


def pd_asof_join():
    result = pd.merge_asof(
        trades_pd_sorted,
        quotes_pd_sorted,
        on="timestamp",
        by="symbol",
        direction="backward",
    )
    force_eval_pandas(result)
    return result

def pl_asof_join():
    """Match trades to most recent quotes via Polars join_asof with backward strategy."""
    result = trades_pl_sorted.join_asof(
        quotes_pl_sorted,
        on="timestamp",
        by="symbol",
        strategy="backward",
    )
    force_eval_polars(result)
    return result


r = benchmark_operation("asof_join", "E_join", pd_asof_join, pl_asof_join)
join_results.append(r)

def pd_anti_join():
    # pandas 3.0 anti-join (or fallback for older versions)
    if IS_PANDAS_3:
        try:
            result = (
                pd.merge(
                    trades_pd_sorted,
                    quotes_pd_sorted[["timestamp", "symbol"]].drop_duplicates(),
                    on=["timestamp", "symbol"],
                    how="left",
                    indicator=True,
                )
                .query("_merge == 'left_only'")
                .drop("_merge", axis=1)
            )
        except Exception:
            # Fallback
            merged = trades_pd_sorted.merge(
                quotes_pd_sorted[["timestamp", "symbol"]].drop_duplicates(),
                on=["timestamp", "symbol"],
                how="left",
                indicator=True,
            )
            result = merged[merged["_merge"] == "left_only"].drop("_merge", axis=1)
    else:
        merged = trades_pd_sorted.merge(
            quotes_pd_sorted[["timestamp", "symbol"]].drop_duplicates(),
            on=["timestamp", "symbol"],
            how="left",
            indicator=True,
        )
        result = merged[merged["_merge"] == "left_only"].drop("_merge", axis=1)
    force_eval_pandas(result)
    return result

def pl_anti_join():
    """Find trades without matching quotes using Polars native anti-join."""
    result = trades_pl_sorted.join(
        quotes_pl_sorted.select(["timestamp", "symbol"]).unique(),
        on=["timestamp", "symbol"],
        how="anti",
    )
    force_eval_polars(result)
    return result


r = benchmark_operation("anti_join", "E_join", pd_anti_join, pl_anti_join)
join_results.append(r)

# Create a smaller lookup table for join benchmark
symbols_df_pd = pd.DataFrame(
    {
        "symbol": ohlcv_pd["symbol"].unique(),
        "sector": np.random.choice(["Tech", "Finance", "Healthcare", "Energy"], size=N_SYMBOLS),
    }
)
symbols_df_pl = pl.DataFrame(
    {
        "symbol": ohlcv_pl.select("symbol").unique()["symbol"],
        "sector": np.random.choice(["Tech", "Finance", "Healthcare", "Energy"], size=N_SYMBOLS),
    }
)


def pd_inner_join():
    result = ohlcv_pd.merge(symbols_df_pd, on="symbol", how="inner")
    force_eval_pandas(result)
    return result

def pl_inner_join():
    """Inner-join OHLCV with sector lookup table using Polars join on symbol."""
    result = ohlcv_pl.join(symbols_df_pl, on="symbol", how="inner")
    force_eval_polars(result)
    return result


r = benchmark_operation("inner_join", "E_join", pd_inner_join, pl_inner_join)
join_results.append(r)

results.extend(join_results)

print("\n" + "=" * 70)
print("CATEGORY F: LAZY/STREAMING")
print("=" * 70)

lazy_results = []

# Save data to parquet for lazy benchmarks
parquet_path = BENCHMARK_DIR / f"ohlcv_{ACTIVE_SCALE.lower()}.parquet"
ohlcv_pl.write_parquet(parquet_path)
print(f"Parquet file: {parquet_path.stat().st_size / 1e6:.1f} MB")

def pd_lazy_filter():
    # pandas: must read all, then filter
    df = pd.read_parquet(parquet_path)
    result = df[df["close"] > price_threshold]
    force_eval_pandas(result)
    return result

def pl_lazy_filter():
    """Scan parquet with predicate pushdown, filtering during read via Polars lazy API."""
    # Polars: predicate pushdown - filter during read
    result = pl.scan_parquet(parquet_path).filter(pl.col("close") > price_threshold).collect()
    force_eval_polars(result)
    return result


r = benchmark_operation("lazy_filter", "F_lazy", pd_lazy_filter, pl_lazy_filter)
lazy_results.append(r)

def pd_column_projection():
    df = pd.read_parquet(parquet_path, columns=["timestamp", "symbol", "close", "volume"])
    force_eval_pandas(df)
    return df

def pl_column_projection():
    """Read only selected columns from parquet using Polars lazy scan with column projection."""
    df = pl.scan_parquet(parquet_path).select(["timestamp", "symbol", "close", "volume"]).collect()
    force_eval_polars(df)
    return df


r = benchmark_operation("column_projection", "F_lazy", pd_column_projection, pl_column_projection)
lazy_results.append(r)

def pd_combined_query():
    df = pd.read_parquet(parquet_path)
    result = (
        df[df["close"] > price_threshold]
        .groupby("symbol")
        .agg({"close": "mean", "volume": "sum"})
        .reset_index()
    )
    force_eval_pandas(result)
    return result

def pl_combined_query():
    """Filter and aggregate in one lazy query, leveraging Polars query optimizer."""
    result = (
        pl.scan_parquet(parquet_path)
        .filter(pl.col("close") > price_threshold)
        .group_by("symbol")
        .agg(
            [
                pl.col("close").mean(),
                pl.col("volume").sum(),
            ]
        )
        .collect()
    )
    force_eval_polars(result)
    return result


r = benchmark_operation("combined_query", "F_lazy", pd_combined_query, pl_combined_query)
lazy_results.append(r)

results.extend(lazy_results)

print("\n" + "=" * 70)
print("CATEGORY G: MEMORY EFFICIENCY")
print("=" * 70)

memory_results = []

process = psutil.Process()

# pandas memory
gc.collect()
mem_before = process.memory_info().rss / 1e6
_ = ohlcv_pd.copy()  # Copy triggers CoW in pandas 3.0
mem_after = process.memory_info().rss / 1e6
pd_mem = mem_after - mem_before

# Polars memory
gc.collect()
mem_before = process.memory_info().rss / 1e6
_ = ohlcv_pl.clone()
mem_after = process.memory_info().rss / 1e6
pl_mem = mem_after - mem_before

print("Copy/Clone operation memory:")
print(f"  pandas: {pd_mem:.1f} MB")
print(f"  Polars: {pl_mem:.1f} MB")

# Estimated DataFrame memory
pd_estimated = ohlcv_pd.memory_usage(deep=True).sum() / 1e6
pl_estimated = ohlcv_pl.estimated_size("mb")

print("\nDataFrame estimated size:")
print(f"  pandas: {pd_estimated:.1f} MB")
print(f"  Polars: {pl_estimated:.1f} MB")
print(f"  Ratio: {pd_estimated / pl_estimated:.2f}x")

memory_results.append(
    {
        "category": "G_memory",
        "operation": "df_estimated_size",
        "pandas_time": pd_estimated,  # Using time fields for memory (MB)
        "polars_time": pl_estimated,
        "speedup": pd_estimated / pl_estimated if pl_estimated > 0 else 1.0,
    }
)

print("\n" + "=" * 70)
print("CATEGORY H: STRING OPERATIONS")
print("=" * 70)

string_results = []

def pd_str_contains():
    result = ohlcv_pd[ohlcv_pd["symbol"].str.contains("SYM_0", regex=False)]
    force_eval_pandas(result)
    return result

def pl_str_contains():
    """Filter rows by literal substring match on symbol using Polars str.contains."""
    result = ohlcv_pl.filter(pl.col("symbol").str.contains("SYM_0", literal=True))
    force_eval_polars(result)
    return result


r = benchmark_operation("str_contains", "H_string", pd_str_contains, pl_str_contains)
string_results.append(r)

def pd_str_replace():
    result = ohlcv_pd.copy()
    result["symbol_new"] = result["symbol"].str.replace("SYM_", "SYMBOL_", regex=False)
    force_eval_pandas(result)
    return result

def pl_str_replace():
    """Replace substring in symbol column using Polars str.replace with literal mode."""
    result = ohlcv_pl.with_columns(
        pl.col("symbol").str.replace("SYM_", "SYMBOL_", literal=True).alias("symbol_new")
    )
    force_eval_polars(result)
    return result


r = benchmark_operation("str_replace", "H_string", pd_str_replace, pl_str_replace)
string_results.append(r)

def pd_str_extract():
    result = ohlcv_pd.copy()
    result["symbol_num"] = result["symbol"].str.extract(r"SYM_(\d+)", expand=False)
    force_eval_pandas(result)
    return result

def pl_str_extract():
    """Extract numeric suffix from symbol via regex capture group using Polars str.extract."""
    result = ohlcv_pl.with_columns(
        pl.col("symbol").str.extract(r"SYM_(\d+)", group_index=1).alias("symbol_num")
    )
    force_eval_polars(result)
    return result


r = benchmark_operation("str_extract", "H_string", pd_str_extract, pl_str_extract)
string_results.append(r)

results.extend(string_results)

print("\n" + "=" * 70)
print("BENCHMARK RESULTS SUMMARY")
print("=" * 70)

# Convert to DataFrame
results_df = pl.DataFrame(results)

# Add memory results if available
if memory_results:
    memory_df = pl.DataFrame(memory_results)
    results_df = pl.concat([results_df, memory_df])

# Summary by category
print("\n### By Category (Mean Speedup)")
category_summary = (
    results_df.group_by("category")
    .agg(
        [
            pl.col("speedup").mean().alias("mean_speedup"),
            pl.col("speedup").min().alias("min_speedup"),
            pl.col("speedup").max().alias("max_speedup"),
            pl.len().alias("n_ops"),
        ]
    )
    .sort("mean_speedup", descending=True)
)
display(category_summary)

# Overall statistics
print("\n### Overall Statistics")
overall_speedup = results_df.select(pl.col("speedup").mean()).item()
print(f"Mean speedup (Polars vs pandas): {overall_speedup:.1f}x")

operations_faster = results_df.filter(pl.col("speedup") > 1.0).height
operations_slower = results_df.filter(pl.col("speedup") < 1.0).height
print(f"Operations where Polars is faster: {operations_faster}/{len(results_df)}")
print(f"Operations where pandas is faster: {operations_slower}/{len(results_df)}")

# Detailed results
print("Detailed Results (sorted by speedup):")
display(results_df.sort("speedup", descending=True))

# Create visualization
fig = make_subplots(
    rows=2,
    cols=2,
    subplot_titles=[
        "Speedup by Category",
        "Operation Times (log scale)",
        "Speedup Distribution",
        "pandas vs Polars Times",
    ],
    specs=[
        [{"type": "bar"}, {"type": "bar"}],
        [{"type": "histogram"}, {"type": "scatter"}],
    ],
)

# 1. Speedup by category (bar chart)
cat_data = category_summary.sort("mean_speedup", descending=True)
fig.add_trace(
    go.Bar(
        x=cat_data["category"].to_list(),
        y=cat_data["mean_speedup"].to_list(),
        marker_color=COLORS["blue"],
        text=[f"{s:.1f}x" for s in cat_data["mean_speedup"].to_list()],
        textposition="outside",
    ),
    row=1,
    col=1,
)
fig.add_hline(y=1.0, line_dash="dash", line_color="gray", row=1, col=1)

# 2. Operation times (grouped bar)
sorted_results = results_df.sort("speedup", descending=True).head(15)
fig.add_trace(
    go.Bar(
        name="pandas",
        x=sorted_results["operation"].to_list(),
        y=sorted_results["pandas_time"].to_list(),
        marker_color=COLORS["amber"],
    ),
    row=1,
    col=2,
)
fig.add_trace(
    go.Bar(
        name="Polars",
        x=sorted_results["operation"].to_list(),
        y=sorted_results["polars_time"].to_list(),
        marker_color=COLORS["blue"],
    ),
    row=1,
    col=2,
)

# 3. Speedup distribution
fig.add_trace(
    go.Histogram(
        x=results_df["speedup"].to_list(),
        nbinsx=20,
        marker_color=COLORS["blue"],
        opacity=0.7,
    ),
    row=2,
    col=1,
)
fig.add_vline(x=1.0, line_dash="dash", line_color="red", row=2, col=1)

# 4. pandas vs Polars scatter
fig.add_trace(
    go.Scatter(
        x=results_df["pandas_time"].to_list(),
        y=results_df["polars_time"].to_list(),
        mode="markers",
        marker=dict(color=COLORS["blue"], size=10),
        text=results_df["operation"].to_list(),
        hovertemplate="%{text}<br>pandas: %{x:.4f}s<br>Polars: %{y:.4f}s<extra></extra>",
    ),
    row=2,
    col=2,
)
# Add diagonal (equal performance line)
max_time = max(results_df["pandas_time"].max(), results_df["polars_time"].max())
fig.add_trace(
    go.Scatter(
        x=[0, max_time],
        y=[0, max_time],
        mode="lines",
        line=dict(dash="dash", color="gray"),
        showlegend=False,
    ),
    row=2,
    col=2,
)

# Update layout
fig.update_xaxes(title_text="Category", row=1, col=1)
fig.update_yaxes(title_text="Speedup (Polars/pandas)", row=1, col=1)

fig.update_xaxes(title_text="Operation", tickangle=45, row=1, col=2)
fig.update_yaxes(title_text="Time (s)", type="log", row=1, col=2)

fig.update_xaxes(title_text="Speedup", row=2, col=1)
fig.update_yaxes(title_text="Count", row=2, col=1)

fig.update_xaxes(title_text="pandas time (s)", row=2, col=2)
fig.update_yaxes(title_text="Polars time (s)", row=2, col=2)

fig.update_layout(
    title_text=f"pandas {PANDAS_VERSION} vs Polars {POLARS_VERSION} Benchmark (Scale: {ACTIVE_SCALE})",
    height=800,
    showlegend=True,
    barmode="group",
)

fig.show()

# Save detailed results
csv_path = RESULTS_DIR / f"pandas_polars_{ACTIVE_SCALE.lower()}.csv"
results_df.write_csv(csv_path)
print(f"Results saved to: {csv_path}")

# Save summary
summary_df = pl.DataFrame(
    {
        "metric": [
            "pandas_version",
            "polars_version",
            "scale",
            "total_rows",
            "mean_speedup",
            "operations_tested",
            "polars_faster_count",
            "pandas_faster_count",
            "cow_enabled",
            "pyarrow_strings",
        ],
        "value": [
            PANDAS_VERSION,
            POLARS_VERSION,
            ACTIVE_SCALE,
            str(total_rows),
            f"{overall_speedup:.2f}",
            str(len(results_df)),
            str(operations_faster),
            str(operations_slower),
            str(COW_ENABLED),
            str(PYARROW_STRINGS),
        ],
    }
)
summary_path = RESULTS_DIR / f"pandas_polars_summary_{ACTIVE_SCALE.lower()}.csv"
summary_df.write_csv(summary_path)
print(f"Summary saved to: {summary_path}")

# Surface this run's category ranking dynamically; the category table and overall
# counts were already displayed in the Summary cell above, so this cell only names
# the top-2 / bottom-2 categories so the takeaways stay in sync with the table.
_ranked = category_summary.sort("mean_speedup", descending=True)
_top2 = _ranked.head(2)["category"].to_list()
_bot2 = _ranked.tail(2)["category"].to_list()
print(f"Scale: {ACTIVE_SCALE} ({total_rows:,} rows)")
print(f"Fastest-on-Polars categories at this scale: {', '.join(_top2)}")
print(f"Smallest-gap (or pandas-faster) categories at this scale: {', '.join(_bot2)}")

print("=" * 70)
print("BENCHMARK COMPLETE")
print("=" * 70)
print(f"pandas {PANDAS_VERSION} vs Polars {POLARS_VERSION}, scale {ACTIVE_SCALE}")
print(f"Results: {csv_path}")