#!/usr/bin/env python3
"""Polars vs. pandas speed comparison.

Reproduces the five tasks from the blog post "Polars vs. pandas: a speed
comparison, through an Apple Silicon lens":

    1. Read   – read_csv / read_parquet with defaults
    2. Write  – write the frame back out with defaults
    3. Compute – add a numeric-expression column
    4. Filter – select rows by condition
    5. Group by – group on a column and aggregate

It runs each task on a small dataset (Iris, CSV) and a large dataset
(Transactions, Parquet), times every run with ``timeit``, and prints the
median over N runs plus the Polars advantage.

Usage
-----
    # Uses the real datasets next to this script if present
    # (scripts/iris.csv and scripts/financial_fraud_detection_dataset.parquet),
    # otherwise synthesises stand-ins – no downloads needed:
    python scripts/polars_vs_pandas_benchmark.py

    # Point at your own Kaggle files explicitly:
    python scripts/polars_vs_pandas_benchmark.py \
        --iris path/to/iris.csv \
        --transactions path/to/transactions.parquet

    # Fewer runs / smaller synthetic data for a quick smoke test:
    python scripts/polars_vs_pandas_benchmark.py --runs 5 --synthetic-rows 1_000_000

Requires: pandas, polars, numpy, pyarrow.
"""

from __future__ import annotations

import argparse
import os
import tempfile
import timeit
from statistics import median

import numpy as np
import pandas as pd
import polars as pl


# --------------------------------------------------------------------------- #
# Dataset preparation
# --------------------------------------------------------------------------- #
def make_iris_csv(path: str) -> None:
    """Write a 150-row, 5-column Iris-shaped CSV (stresses start-up overhead)."""
    rng = np.random.default_rng(0)
    n = 150
    species = np.repeat(["setosa", "versicolor", "virginica"], n // 3)
    df = pd.DataFrame(
        {
            "sepal_length": rng.normal(5.8, 0.8, n).round(1),
            "sepal_width": rng.normal(3.0, 0.4, n).round(1),
            "petal_length": rng.normal(3.8, 1.8, n).round(1),
            "petal_width": rng.normal(1.2, 0.8, n).round(1),
            "species": species,
        }
    )
    df.to_csv(path, index=False)


def make_transactions_parquet(path: str, rows: int, compression: str = "zstd") -> None:
    """Write a wide (24-col) transactions Parquet file (stresses scalability)."""
    rng = np.random.default_rng(1)
    data = {
        "transaction_id": np.arange(rows, dtype=np.int64),
        "account_id": rng.integers(0, 250_000, rows, dtype=np.int64),
        "merchant_id": rng.integers(0, 50_000, rows, dtype=np.int64),
        "category": rng.integers(0, 40, rows).astype("U2"),
        "amount": rng.normal(120.0, 480.0, rows).round(2),
        "balance": rng.normal(5000.0, 9000.0, rows).round(2),
        "is_fraud": rng.integers(0, 2, rows, dtype=np.int8),
    }
    # Pad out to ~24 columns with assorted numeric/string fields.
    for i in range(6):
        data[f"num_{i}"] = rng.normal(0.0, 1.0, rows)
    for i in range(6):
        data[f"flag_{i}"] = rng.integers(0, 2, rows, dtype=np.int8)
    for i in range(5):
        data[f"code_{i}"] = rng.integers(0, 1000, rows).astype("U3")
    pd.DataFrame(data).to_parquet(path, index=False, compression=compression)


# --------------------------------------------------------------------------- #
# Timing helper
# --------------------------------------------------------------------------- #
def bench(stmt, runs: int) -> float:
    """Median wall-clock seconds of `stmt` over `runs` single executions."""
    times = timeit.repeat(stmt, repeat=runs, number=1)
    return median(times)


def fmt(seconds: float) -> str:
    """Human-friendly duration, matching the article's ms / s style."""
    if seconds < 1.0:
        return f"{seconds * 1e3:.2f} ms"
    return f"{seconds:.2f} s"


# --------------------------------------------------------------------------- #
# Task definitions
# --------------------------------------------------------------------------- #
def pick_numeric_col(df: pd.DataFrame, preferred: str | None) -> str:
    """Return `preferred` if it's a usable numeric column, else the first numeric one."""
    if preferred and preferred in df.columns and pd.api.types.is_numeric_dtype(df[preferred]):
        return preferred
    for col in df.columns:
        if pd.api.types.is_numeric_dtype(df[col]):
            return col
    raise ValueError("No numeric column found for the compute/filter/group-by tasks.")


def pick_group_col(df: pd.DataFrame, preferred: str | None, exclude: str) -> str:
    """Return `preferred` if present, else the lowest-cardinality non-numeric column.

    Falls back to the column with the fewest distinct values so the group-by is
    meaningful (avoids grouping by a unique id).
    """
    if preferred and preferred in df.columns:
        return preferred
    candidates = [c for c in df.columns if c != exclude and not pd.api.types.is_numeric_dtype(df[c])]
    pool = candidates or [c for c in df.columns if c != exclude]
    return min(pool, key=lambda c: df[c].nunique())


def run_suite(
    name: str,
    csv_path: str,
    parquet_path: str,
    runs: int,
    is_csv: bool,
    num_col: str | None = None,
    grp_col: str | None = None,
    compression: str = "zstd",
):
    """Run the five tasks on one dataset for both libraries and print a table."""
    src = csv_path if is_csv else parquet_path
    reader = "read_csv" if is_csv else "read_parquet"

    # Preload frames for the compute/filter/group-by tasks.
    pd_df = pd.read_csv(src) if is_csv else pd.read_parquet(src)
    pl_df = pl.read_csv(src) if is_csv else pl.read_parquet(src)

    # Resolve which numeric / group columns to use: CLI override → preferred
    # default for the synthetic data → auto-detected from the real file's schema.
    preferred_num = num_col or ("petal_length" if is_csv else "amount")
    preferred_grp = grp_col or ("species" if is_csv else "merchant_category")
    num_col = pick_numeric_col(pd_df, preferred_num)
    grp_col = pick_group_col(pd_df, preferred_grp, exclude=num_col)
    print(
        f"  rows={len(pd_df):,}  cols={len(pd_df.columns)}  "
        f"numeric='{num_col}'  group='{grp_col}'"
    )

    out_dir = tempfile.mkdtemp(prefix="pvp_")
    pd_out = os.path.join(out_dir, "out_pd" + (".csv" if is_csv else ".parquet"))
    pl_out = os.path.join(out_dir, "out_pl" + (".csv" if is_csv else ".parquet"))

    # --- pandas statements ------------------------------------------------- #
    if is_csv:
        pd_read = lambda: pd.read_csv(src)
        pd_write = lambda: pd_df.to_csv(pd_out, index=False)
        pl_read = lambda: pl.read_csv(src)
        pl_write = lambda: pl_df.write_csv(pl_out)
    else:
        pd_read = lambda: pd.read_parquet(src)
        pd_write = lambda: pd_df.to_parquet(pd_out, index=False, compression=compression)
        pl_read = lambda: pl.read_parquet(src)
        pl_write = lambda: pl_df.write_parquet(pl_out, compression=compression)

    pd_compute = lambda: pd_df.assign(_new=pd_df[num_col] * 2.0 + 1.0)
    pd_filter = lambda: pd_df[pd_df[num_col] > pd_df[num_col].mean()]
    pd_group = lambda: pd_df.groupby(grp_col)[num_col].mean()

    pl_compute = lambda: pl_df.with_columns((pl.col(num_col) * 2.0 + 1.0).alias("_new"))
    pl_filter = lambda: pl_df.filter(pl.col(num_col) > pl.col(num_col).mean())
    pl_group = lambda: pl_df.group_by(grp_col).agg(pl.col(num_col).mean())

    tasks = [
        (f"Read ({reader})", pd_read, pl_read),
        ("Write", pd_write, pl_write),
        ("Numeric expr", pd_compute, pl_compute),
        ("Filter", pd_filter, pl_filter),
        ("Group by", pd_group, pl_group),
    ]

    print(f"\n=== {name} ===")
    print(f"{'Task':<22}{'pandas':>12}{'Polars':>12}{'Advantage':>12}")
    print("-" * 58)
    for label, pd_stmt, pl_stmt in tasks:
        t_pd = bench(pd_stmt, runs)
        t_pl = bench(pl_stmt, runs)
        adv = t_pd / t_pl if t_pl > 0 else float("inf")
        print(f"{label:<22}{fmt(t_pd):>12}{fmt(t_pl):>12}{adv:>11.1f}x")


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #
def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--iris", help="Path to an Iris-style CSV (small dataset).")
    parser.add_argument(
        "--transactions", help="Path to a Transactions-style Parquet (large dataset)."
    )
    parser.add_argument(
        "--runs", type=int, default=10, help="Timed runs per task (median reported)."
    )
    parser.add_argument(
        "--synthetic-rows",
        type=int,
        default=7_480_000,
        help="Rows for the synthetic Transactions file when --transactions is omitted.",
    )
    parser.add_argument(
        "--tx-num-col",
        help="Numeric column to compute/filter/group on in the Transactions file "
        "(auto-detected if omitted or absent).",
    )
    parser.add_argument(
        "--tx-group-col",
        help="Column to group by in the Transactions file "
        "(auto-detected, lowest-cardinality, if omitted or absent).",
    )
    parser.add_argument(
        "--compression",
        default="zstd",
        help="Parquet compression codec for the Write task and synthetic file "
        "(default: zstd; e.g. snappy, gzip, lz4, none).",
    )
    args = parser.parse_args()

    print(f"pandas {pd.__version__}  ·  polars {pl.__version__}  ·  numpy {np.__version__}")
    print(f"Timing: median of {args.runs} runs via timeit  ·  Parquet codec: {args.compression}")

    # Use the real datasets sitting next to this script if present; otherwise
    # synthesise stand-ins. Explicit --iris / --transactions always win.
    here = os.path.dirname(os.path.abspath(__file__))
    default_iris = os.path.join(here, "iris.csv")
    default_tx = os.path.join(here, "financial_fraud_detection_dataset.parquet")
    tmp = tempfile.mkdtemp(prefix="pvp_data_")

    iris_path = args.iris or (default_iris if os.path.exists(default_iris) else None)
    if iris_path:
        print(f"Iris dataset: {iris_path}")
    else:
        iris_path = os.path.join(tmp, "iris.csv")
        print("No Iris file found – synthesising Iris CSV (150 rows)...")
        make_iris_csv(iris_path)

    tx_path = args.transactions or (default_tx if os.path.exists(default_tx) else None)
    if tx_path:
        print(f"Transactions dataset: {tx_path}")
    else:
        tx_path = os.path.join(tmp, "transactions.parquet")
        print(f"No Transactions file found – synthesising Parquet ({args.synthetic_rows:,} rows)...")
        make_transactions_parquet(tx_path, args.synthetic_rows, compression=args.compression)

    run_suite("Small / CSV (Iris)", iris_path, "", args.runs, is_csv=True)
    run_suite(
        "Large / Parquet (Transactions)",
        "",
        tx_path,
        args.runs,
        is_csv=False,
        num_col=args.tx_num_col,
        grp_col=args.tx_group_col,
        compression=args.compression,
    )


if __name__ == "__main__":
    main()