#!/usr/bin/env python3
"""CSV vs. Parquet: a storage and performance comparison.

Reproduces the measurements from the blog post "CSV vs. Parquet: which file
format should you choose for modern data platforms?". It runs two things:

1. A head-to-head between CSV and Parquet on a small dataset (the classic
   Iris dataset, the kind of tiny file that pours out of ERP and Excel
   exports) – size on disk, full read, write, single-column read and a
   filtered scan.

2. A Parquet *compression sweep* on a large dataset (by default a real
   5-million-row financial-fraud dataset), comparing CSV against every
   Parquet codec – none, snappy, lz4, zstd, gzip, brotli – across storage
   size, write, full read, single-column read, filtered read and a
   projected group-by. This is where you see how much the codec choice
   actually buys you.

Timings use ``timeit``. Fast operations are timed over ``--runs`` executions;
operations that take more than a couple of seconds (the big CSV reads/writes)
fall back to fewer runs automatically so the whole thing finishes in minutes.

Usage
-----
    # Uses scripts/financial_fraud_detection_dataset.csv if present, else
    # synthesises a stand-in large dataset. Iris is always synthesised.
    python scripts/csv_vs_parquet_benchmark.py

    # Point at any large CSV for the compression sweep:
    python scripts/csv_vs_parquet_benchmark.py --large path/to/data.csv

    # Quicker pass on a sample of the large file:
    python scripts/csv_vs_parquet_benchmark.py --sample-rows 1_000_000 --runs 5

Requires: pandas, pyarrow, numpy.
"""

from __future__ import annotations

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

import numpy as np
import pandas as pd
import pyarrow

# Default to the real dataset shipped alongside this script, if present.
HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_LARGE_CSV = os.path.join(HERE, "financial_fraud_detection_dataset.csv")

# Parquet codecs to sweep. None == uncompressed. Order = low to high effort.
PARQUET_CODECS = [None, "snappy", "lz4", "zstd", "gzip", "brotli"]


# --------------------------------------------------------------------------- #
# Dataset preparation
# --------------------------------------------------------------------------- #
def make_iris(rng: np.random.Generator) -> pd.DataFrame:
    """A 150-row, 5-column Iris-shaped frame (the small / interoperability case)."""
    n = 150
    species = np.repeat(["setosa", "versicolor", "virginica"], n // 3)
    return 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,
        }
    )


def make_transactions(rng: np.random.Generator, rows: int) -> pd.DataFrame:
    """A wide (18-col) transactions frame used when no real CSV is supplied."""
    countries = np.array(["Tokyo", "Toronto", "London", "Berlin", "Sydney", "Paris"])
    categories = np.array(
        ["groceries", "fuel", "travel", "utilities", "retail", "online",
         "entertainment", "restaurant"]
    )
    channels = np.array(["card", "ACH", "wire", "UPI"])
    devices = np.array(["mobile", "atm", "web", "pos"])
    data = {
        "transaction_id": np.arange(rows, dtype=np.int64),
        "sender_account": rng.integers(0, 250_000, rows, dtype=np.int64),
        "receiver_account": rng.integers(0, 250_000, rows, dtype=np.int64),
        "amount": rng.normal(400.0, 300.0, rows).round(2),
        "transaction_type": rng.choice(["withdrawal", "deposit", "transfer"], rows),
        "merchant_category": categories[rng.integers(0, len(categories), rows)],
        "location": countries[rng.integers(0, len(countries), rows)],
        "device_used": devices[rng.integers(0, len(devices), rows)],
        "is_fraud": rng.integers(0, 2, rows, dtype=np.int8).astype(bool),
        "spending_deviation_score": rng.normal(0.0, 1.0, rows).round(2),
        "velocity_score": rng.integers(0, 20, rows, dtype=np.int64),
        "geo_anomaly_score": rng.random(rows).round(2),
        "payment_channel": channels[rng.integers(0, len(channels), rows)],
    }
    return pd.DataFrame(data)


def load_large(path: str | None, sample_rows: int | None,
               rng: np.random.Generator) -> tuple[pd.DataFrame, str]:
    """Load the large dataset from `path`, or synthesise one if unavailable."""
    if path and os.path.exists(path):
        label = os.path.basename(path)
        print(f"Loading large dataset from {label} ...")
        df = pd.read_csv(path, nrows=sample_rows)
        return df, f"Real: {label}"
    rows = sample_rows or 5_000_000
    print(f"No large CSV found – synthesising {rows:,} transactions ...")
    return make_transactions(rng, rows), "Synthetic transactions"


# --------------------------------------------------------------------------- #
# Timing helpers
# --------------------------------------------------------------------------- #
def bench(stmt, runs: int) -> float:
    """Median wall-clock seconds of `stmt`, with adaptive run counts.

    Slow operations (big CSV reads/writes) are capped at a few runs so the
    suite stays quick; sub-second operations get the full `runs` count.
    """
    first = timeit.timeit(stmt, number=1)
    if first >= 2.0:
        n_more = max(0, min(runs, 3) - 1)
    elif first >= 0.2:
        n_more = max(0, min(runs, 5) - 1)
    else:
        n_more = runs - 1
    times = [first]
    if n_more:
        times += timeit.repeat(stmt, repeat=n_more, number=1)
    return median(times)


def fmt_time(seconds: float) -> str:
    if seconds < 1.0:
        return f"{seconds * 1e3:.1f} ms"
    return f"{seconds:.2f} s"


def fmt_size(num_bytes: int) -> str:
    size = float(num_bytes)
    for unit in ["B", "KB", "MB", "GB"]:
        if size < 1024 or unit == "GB":
            return f"{size:.1f} {unit}"
        size /= 1024
    return f"{size:.1f} GB"


def gzip_file(src: str, dst: str) -> None:
    with open(src, "rb") as f_in, gzip.open(dst, "wb", compresslevel=6) as f_out:
        shutil.copyfileobj(f_in, f_out)


# --------------------------------------------------------------------------- #
# Suite 1: CSV vs Parquet head-to-head
# --------------------------------------------------------------------------- #
def run_suite(name: str, df: pd.DataFrame, num_col: str, proj_col: str,
              runs: int, out_dir: str) -> None:
    """Materialise both formats, then time four comparisons and print a table."""
    csv_path = os.path.join(out_dir, "data.csv")
    gz_path = os.path.join(out_dir, "data.csv.gz")
    pq_path = os.path.join(out_dir, "data.parquet")

    df.to_csv(csv_path, index=False)
    df.to_parquet(pq_path, index=False, compression="snappy")
    gzip_file(csv_path, gz_path)

    csv_size = os.path.getsize(csv_path)
    gz_size = os.path.getsize(gz_path)
    pq_size = os.path.getsize(pq_path)

    print(f"\n=== {name} ({len(df):,} rows x {df.shape[1]} cols) ===")
    print("Storage on disk")
    print(f"  {'CSV':<16}{fmt_size(csv_size):>12}")
    print(f"  {'CSV (gzip)':<16}{fmt_size(gz_size):>12}   "
          f"{csv_size / gz_size:.1f}x smaller than CSV")
    print(f"  {'Parquet':<16}{fmt_size(pq_size):>12}   "
          f"{csv_size / pq_size:.1f}x smaller than CSV")

    csv_read = lambda: pd.read_csv(csv_path)
    pq_read = lambda: pd.read_parquet(pq_path)
    csv_write = lambda: df.to_csv(os.path.join(out_dir, "w.csv"), index=False)
    pq_write = lambda: df.to_parquet(
        os.path.join(out_dir, "w.parquet"), index=False, compression="snappy"
    )
    csv_project = lambda: pd.read_csv(csv_path, usecols=[proj_col])
    pq_project = lambda: pd.read_parquet(pq_path, columns=[proj_col])

    def csv_filter():
        d = pd.read_csv(csv_path)
        return d[d[num_col] > d[num_col].mean()]

    def pq_filter():
        d = pd.read_parquet(pq_path)
        return d[d[num_col] > d[num_col].mean()]

    tasks = [
        ("Full read", csv_read, pq_read),
        ("Write", csv_write, pq_write),
        (f"Read 1 column ({proj_col})", csv_project, pq_project),
        ("Read + filter", csv_filter, pq_filter),
    ]

    print(f"\nTiming (median of {runs} runs)")
    print(f"  {'Task':<26}{'CSV':>12}{'Parquet':>12}{'Speed-up':>12}")
    print("  " + "-" * 60)
    for label, csv_stmt, pq_stmt in tasks:
        t_csv = bench(csv_stmt, runs)
        t_pq = bench(pq_stmt, runs)
        speedup = t_csv / t_pq if t_pq > 0 else float("inf")
        print(f"  {label:<26}{fmt_time(t_csv):>12}{fmt_time(t_pq):>12}"
              f"{speedup:>11.1f}x")


# --------------------------------------------------------------------------- #
# Suite 2: Parquet compression sweep
# --------------------------------------------------------------------------- #
def run_compression_sweep(name: str, df: pd.DataFrame, num_col: str,
                          grp_col: str, runs: int, out_dir: str) -> None:
    """Compare CSV against every Parquet codec on storage and read operations."""
    proj_cols = [num_col]
    grp_proj = [grp_col, num_col]

    # --- Materialise every format once, recording size + write time -------- #
    csv_path = os.path.join(out_dir, "sweep.csv")
    gz_path = os.path.join(out_dir, "sweep.csv.gz")
    df.to_csv(csv_path, index=False)
    gzip_file(csv_path, gz_path)

    formats: list[tuple[str, str]] = [("CSV", csv_path)]  # (label, path)
    write_times: dict[str, float] = {}

    print(f"\n=== Parquet compression sweep: {name} "
          f"({len(df):,} rows x {df.shape[1]} cols) ===")

    # CSV write time (for reference).
    write_times["CSV"] = bench(
        lambda: df.to_csv(os.path.join(out_dir, "w.csv"), index=False), runs
    )

    for codec in PARQUET_CODECS:
        label = f"Parquet ({codec or 'none'})"
        path = os.path.join(out_dir, f"sweep_{codec or 'none'}.parquet")
        write_times[label] = bench(
            lambda p=path, c=codec: df.to_parquet(p, index=False, compression=c),
            runs,
        )
        formats.append((label, path))

    # --- Storage table ----------------------------------------------------- #
    csv_size = os.path.getsize(csv_path)
    print("\nStorage on disk")
    print(f"  {'Format':<20}{'Size':>12}{'vs CSV':>12}")
    print("  " + "-" * 44)
    print(f"  {'CSV':<20}{fmt_size(csv_size):>12}{'1.0x':>12}")
    gz_size = os.path.getsize(gz_path)
    print(f"  {'CSV (gzip)':<20}{fmt_size(gz_size):>12}"
          f"{csv_size / gz_size:>11.1f}x")
    for label, path in formats[1:]:
        size = os.path.getsize(path)
        print(f"  {label:<20}{fmt_size(size):>12}{csv_size / size:>11.1f}x")

    # --- Read/operation timing table --------------------------------------- #
    def read_fn(label, path):
        return (lambda: pd.read_csv(path)) if label == "CSV" \
            else (lambda: pd.read_parquet(path))

    def project_fn(label, path):
        return (lambda: pd.read_csv(path, usecols=proj_cols)) if label == "CSV" \
            else (lambda: pd.read_parquet(path, columns=proj_cols))

    def filter_fn(label, path):
        if label == "CSV":
            def f():
                d = pd.read_csv(path)
                return d[d[num_col] > d[num_col].mean()]
        else:
            def f():
                d = pd.read_parquet(path)
                return d[d[num_col] > d[num_col].mean()]
        return f

    def groupby_fn(label, path):
        if label == "CSV":
            def f():
                d = pd.read_csv(path, usecols=grp_proj)
                return d.groupby(grp_col)[num_col].sum()
        else:
            def f():
                d = pd.read_parquet(path, columns=grp_proj)
                return d.groupby(grp_col)[num_col].sum()
        return f

    print(f"\nTiming (median, adaptive runs up to {runs})")
    print(f"  {'Format':<20}{'Write':>10}{'Read':>10}"
          f"{'1-col':>10}{'Filter':>10}{'GroupBy':>10}")
    print("  " + "-" * 70)
    for label, path in formats:
        t_read = bench(read_fn(label, path), runs)
        t_proj = bench(project_fn(label, path), runs)
        t_filt = bench(filter_fn(label, path), runs)
        t_grp = bench(groupby_fn(label, path), runs)
        print(f"  {label:<20}{fmt_time(write_times[label]):>10}"
              f"{fmt_time(t_read):>10}{fmt_time(t_proj):>10}"
              f"{fmt_time(t_filt):>10}{fmt_time(t_grp):>10}")


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #
def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--iris", help="Path to a real Iris-style CSV (small dataset).")
    parser.add_argument(
        "--large", default=DEFAULT_LARGE_CSV,
        help="Path to the large CSV for the compression sweep "
             "(defaults to scripts/financial_fraud_detection_dataset.csv).",
    )
    parser.add_argument("--runs", type=int, default=7,
                        help="Max timed runs per task (median reported).")
    parser.add_argument("--sample-rows", type=int, default=None,
                        help="Cap the large dataset to this many rows.")
    args = parser.parse_args()

    print(f"pandas {pd.__version__}  ·  pyarrow {pyarrow.__version__}"
          f"  ·  numpy {np.__version__}")
    print(f"Timing: median of up to {args.runs} runs via timeit (adaptive)")

    rng = np.random.default_rng(42)

    # Small dataset (Iris).
    if args.iris:
        iris = pd.read_csv(args.iris)
        num = "petal_length" if "petal_length" in iris.columns \
            else iris.select_dtypes("number").columns[0]
        proj = "sepal_length" if "sepal_length" in iris.columns else num
    else:
        print("Synthesising Iris (150 rows)...")
        iris = make_iris(rng)
        num, proj = "petal_length", "sepal_length"

    # Large dataset (real fraud CSV by default).
    large, large_label = load_large(args.large, args.sample_rows, rng)
    large_num = "amount" if "amount" in large.columns else \
        large.select_dtypes("number").columns[0]
    large_grp = "merchant_category" if "merchant_category" in large.columns else \
        large.select_dtypes("object").columns[0]

    out = tempfile.mkdtemp(prefix="cvp_")
    try:
        run_suite("Small / Iris", iris, num, proj, args.runs, out)
        run_compression_sweep(large_label, large, large_num, large_grp,
                               args.runs, out)
    finally:
        shutil.rmtree(out, ignore_errors=True)


if __name__ == "__main__":
    main()