#!/usr/bin/env python3
"""
Compare performance of robin-sparkless vs sparkless (PySpark drop-in) on the same pipelines.

Requirements:
  - robin_sparkless: build with `maturin develop --features pyo3` from repo root
  - sparkless: pip install sparkless

Usage:
  python scripts/bench_robin_vs_sparkless.py
  python scripts/bench_robin_vs_sparkless.py --quick          # smaller sizes, fewer iterations
  python scripts/bench_robin_vs_sparkless.py --sizes 100000 500000 1000000
  python scripts/bench_robin_vs_sparkless.py --check-results

Output: Table of mean time (and std) per backend per pipeline; speedup(robin) when both available.
"""

import argparse
import sys
import time
import warnings
from typing import Any, Callable, List, Optional, Tuple

# Suppress Sparkless Polars LazyFrame schema warnings
warnings.filterwarnings("ignore", message=".*LazyFrame.*")
warnings.filterwarnings("ignore", message=".*schema.*", module=".*materializer.*")

# Defaults: realistic and big
DEFAULT_SIZES = [100_000, 500_000]
QUICK_SIZES = [10_000, 50_000]
DEFAULT_ITERATIONS = 4
QUICK_ITERATIONS = 2
WARMUP_ITERATIONS = 1

# Join/union pipeline sizes (fixed)
JOIN_SIZE = 50_000  # each side
UNION_SIZE = 200_000  # each df, then union
TOP_N = 10_000


def make_data(n: int) -> List[Tuple[int, int, str]]:
    """(id, age, name) - age 0..79, names user_0, user_1, ..."""
    return [(i, i % 80, f"user_{i}") for i in range(n)]


# ---------- Pipeline 1: filter -> select -> groupBy count ----------
def run_robin_filter_select_groupby(data: List[Tuple[int, int, str]], cols: List[str]) -> Any:
    import robin_sparkless as rs
    spark = rs.SparkSession.builder().app_name("bench").get_or_create()
    df = spark.create_dataframe(data, cols)
    return df.filter(rs.col("age").gt(rs.lit(30))).select(["name", "age"]).group_by(["age"]).count().collect()


def run_sparkless_filter_select_groupby(data: List[Tuple[int, int, str]], cols: List[str]) -> Any:
    from sparkless.sql import SparkSession
    import sparkless.sql.functions as F
    spark = SparkSession.builder.appName("bench").getOrCreate()
    df = spark.createDataFrame(data, cols)
    return df.filter(F.col("age") > F.lit(30)).select("name", "age").groupBy("age").count().collect()


# ---------- Pipeline 2: groupBy(age).agg(sum(id), count) ----------
def run_robin_groupby_multi_agg(data: List[Tuple[int, int, str]], cols: List[str]) -> Any:
    import robin_sparkless as rs
    spark = rs.SparkSession.builder().app_name("bench").get_or_create()
    df = spark.create_dataframe(data, cols)
    grouped = df.group_by(["age"])
    return grouped.agg([
        rs.sum(rs.col("id")).alias("sum_id"),
        rs.count(rs.col("id")).alias("count_id"),
    ]).collect()


def run_sparkless_groupby_multi_agg(data: List[Tuple[int, int, str]], cols: List[str]) -> Any:
    from sparkless.sql import SparkSession
    import sparkless.sql.functions as F
    spark = SparkSession.builder.appName("bench").getOrCreate()
    df = spark.createDataFrame(data, cols)
    return df.groupBy("age").agg(F.sum("id").alias("sum_id"), F.count("id").alias("count_id")).collect()


# ---------- Pipeline 3: orderBy(age).limit(TOP_N) ----------
def run_robin_order_by_limit(data: List[Tuple[int, int, str]], cols: List[str], n: int) -> Any:
    import robin_sparkless as rs
    spark = rs.SparkSession.builder().app_name("bench").get_or_create()
    df = spark.create_dataframe(data, cols)
    return df.order_by(["age"], ascending=[False]).limit(n).collect()


def run_sparkless_order_by_limit(data: List[Tuple[int, int, str]], cols: List[str], n: int) -> Any:
    from sparkless.sql import SparkSession
    import sparkless.sql.functions as F
    spark = SparkSession.builder.appName("bench").getOrCreate()
    df = spark.createDataFrame(data, cols)
    return df.orderBy(F.desc("age")).limit(n).collect()


# ---------- Pipeline 4: distinct ----------
def run_robin_distinct(data: List[Tuple[int, int, str]], cols: List[str]) -> Any:
    import robin_sparkless as rs
    spark = rs.SparkSession.builder().app_name("bench").get_or_create()
    df = spark.create_dataframe(data, cols)
    return df.distinct(None).collect()


def run_sparkless_distinct(data: List[Tuple[int, int, str]], cols: List[str]) -> Any:
    from sparkless.sql import SparkSession
    spark = SparkSession.builder.appName("bench").getOrCreate()
    df = spark.createDataFrame(data, cols)
    return df.distinct().collect()


# ---------- Pipeline 5: union two DataFrames then groupBy count ----------
def run_robin_union_then_groupby(
    data_a: List[Tuple[int, int, str]],
    data_b: List[Tuple[int, int, str]],
    cols: List[str],
) -> Any:
    import robin_sparkless as rs
    spark = rs.SparkSession.builder().app_name("bench").get_or_create()
    df_a = spark.create_dataframe(data_a, cols)
    df_b = spark.create_dataframe(data_b, cols)
    return df_a.union(df_b).group_by(["age"]).count().collect()


def run_sparkless_union_then_groupby(
    data_a: List[Tuple[int, int, str]],
    data_b: List[Tuple[int, int, str]],
    cols: List[str],
) -> Any:
    from sparkless.sql import SparkSession
    spark = SparkSession.builder.appName("bench").getOrCreate()
    df_a = spark.createDataFrame(data_a, cols)
    df_b = spark.createDataFrame(data_b, cols)
    # unionByName so column order doesn't matter (Sparkless may reorder)
    return df_a.unionByName(df_b).groupBy("age").count().collect()


# ---------- Pipeline 6: inner join (JOIN_SIZE x JOIN_SIZE) ----------
def run_robin_join(data_a: List[Tuple[int, int, str]], data_b: List[Tuple[int, int, str]]) -> Any:
    import robin_sparkless as rs
    spark = rs.SparkSession.builder().app_name("bench").get_or_create()
    df_a = spark.create_dataframe(data_a, ["id", "age", "name"])
    df_b = spark.create_dataframe(data_b, ["id", "x", "y"])
    return df_a.join(df_b, ["id"], "inner").collect()


def run_sparkless_join(data_a: List[Tuple[int, int, str]], data_b: List[Tuple[int, int, str]]) -> Any:
    from sparkless.sql import SparkSession
    spark = SparkSession.builder.appName("bench").getOrCreate()
    df_a = spark.createDataFrame(data_a, ["id", "age", "name"])
    df_b = spark.createDataFrame(data_b, ["id", "x", "y"])
    return df_a.join(df_b, "id", "inner").collect()


def time_it(
    fn: Callable[[], Any],
    iterations: int,
    warmup: int = WARMUP_ITERATIONS,
) -> Tuple[float, float]:
    for _ in range(warmup):
        fn()
    times: List[float] = []
    for _ in range(iterations):
        t0 = time.perf_counter()
        fn()
        t1 = time.perf_counter()
        times.append(t1 - t0)
    n = len(times)
    mean = sum(times) / n
    variance = sum((t - mean) ** 2 for t in times) / n if n > 1 else 0.0
    std = variance ** 0.5
    return mean, std


def run_table(
    name: str,
    sizes: List[int],
    has_robin: bool,
    has_sparkless: bool,
    iterations: int,
    check_results: bool,
    make_row: Callable[[int], Tuple[Optional[float], Optional[float], Optional[float], Optional[float], Any]],
) -> None:
    """Print a benchmark table. make_row(n) returns (robin_mean, robin_std, sparkless_mean, sparkless_std, match?)."""
    print()
    print(name)
    print("-" * 72)
    header = ["n_rows", "robin_mean(s)", "robin_std(s)", "sparkless_mean(s)", "sparkless_std(s)", "speedup(robin)"]
    if not has_robin:
        header = ["n_rows", "sparkless_mean(s)", "sparkless_std(s)"]
    if not has_sparkless:
        header = ["n_rows", "robin_mean(s)", "robin_std(s)"]
    print("  ".join(f"{h:>18}" for h in header))

    for n in sizes:
        robin_mean, robin_std, sparkless_mean, sparkless_std, _ = make_row(n)

        parts = [f"{n:>18}"]
        if has_robin:
            parts.append(f"{robin_mean:>18.4f}" if robin_mean is not None else " " * 18)
            parts.append(f"{robin_std:>18.4f}" if robin_std is not None else " " * 18)
        if has_sparkless:
            parts.append(f"{sparkless_mean:>18.4f}" if sparkless_mean is not None else " " * 18)
            parts.append(f"{sparkless_std:>18.4f}" if sparkless_std is not None else " " * 18)
        if has_robin and has_sparkless and robin_mean and sparkless_mean and sparkless_mean > 0:
            parts.append(f"{(sparkless_mean / robin_mean):>18.2f}x")
        print("  ".join(parts))


def main() -> int:
    ap = argparse.ArgumentParser(description="Benchmark robin-sparkless vs sparkless")
    ap.add_argument("--sizes", type=int, nargs="+", default=None, help="Row counts (default: 100k, 500k; use --quick for 10k, 50k)")
    ap.add_argument("--iterations", type=int, default=None, help="Timing iterations (default: 4; 2 with --quick)")
    ap.add_argument("--quick", action="store_true", help="Smaller sizes and fewer iterations")
    ap.add_argument("--check-results", action="store_true", help="Verify result equality when both backends available")
    ap.add_argument("--no-sparkless", action="store_true", help="Skip sparkless")
    ap.add_argument("--no-robin", action="store_true", help="Skip robin-sparkless")
    args = ap.parse_args()

    sizes = args.sizes if args.sizes is not None else (QUICK_SIZES if args.quick else DEFAULT_SIZES)
    iterations = args.iterations if args.iterations is not None else (QUICK_ITERATIONS if args.quick else DEFAULT_ITERATIONS)

    has_robin = False
    has_sparkless = False
    if not args.no_robin:
        try:
            import robin_sparkless  # noqa: F401
            has_robin = True
        except ImportError as e:
            print("robin_sparkless not available:", e, file=sys.stderr)
    if not args.no_sparkless:
        try:
            from sparkless.sql import SparkSession  # noqa: F401
            has_sparkless = True
        except ImportError as e:
            print("sparkless not available:", e, file=sys.stderr)

    if not has_robin and not has_sparkless:
        print("Need at least one of robin_sparkless or sparkless.", file=sys.stderr)
        return 1

    # Suppress Sparkless LazyFrame schema warnings for cleaner output
    cols = ["id", "age", "name"]

    # ----- Pipeline 1: filter -> select -> groupBy count -----
    def row1(n: int):
        data = make_data(n)
        r_m, r_s = (time_it(lambda: run_robin_filter_select_groupby(data, cols), iterations) if has_robin else (None, None))
        s_m, s_s = (time_it(lambda: run_sparkless_filter_select_groupby(data, cols), iterations) if has_sparkless else (None, None))
        return r_m, r_s, s_m, s_s, None

    run_table(
        "Pipeline 1: filter(age>30) -> select -> groupBy(age).count()",
        sizes,
        has_robin,
        has_sparkless,
        iterations,
        args.check_results,
        row1,
    )

    # ----- Pipeline 2: groupBy(age).agg(sum(id), count(id)) -----
    def row2(n: int):
        data = make_data(n)
        r_m, r_s = (time_it(lambda: run_robin_groupby_multi_agg(data, cols), iterations) if has_robin else (None, None))
        s_m, s_s = (time_it(lambda: run_sparkless_groupby_multi_agg(data, cols), iterations) if has_sparkless else (None, None))
        return r_m, r_s, s_m, s_s, None

    run_table(
        "Pipeline 2: groupBy(age).agg(sum(id), count(id))",
        sizes,
        has_robin,
        has_sparkless,
        iterations,
        args.check_results,
        row2,
    )

    # ----- Pipeline 3: orderBy(age desc).limit(TOP_N) -----
    def row3(n: int):
        data = make_data(n)
        r_m, r_s = (time_it(lambda: run_robin_order_by_limit(data, cols, TOP_N), iterations) if has_robin else (None, None))
        s_m, s_s = (time_it(lambda: run_sparkless_order_by_limit(data, cols, TOP_N), iterations) if has_sparkless else (None, None))
        return r_m, r_s, s_m, s_s, None

    run_table(
        f"Pipeline 3: orderBy(age desc).limit({TOP_N})",
        sizes,
        has_robin,
        has_sparkless,
        iterations,
        args.check_results,
        row3,
    )

    # ----- Pipeline 4: distinct -----
    def row4(n: int):
        data = make_data(n)
        r_m, r_s = (time_it(lambda: run_robin_distinct(data, cols), iterations) if has_robin else (None, None))
        s_m, s_s = (time_it(lambda: run_sparkless_distinct(data, cols), iterations) if has_sparkless else (None, None))
        return r_m, r_s, s_m, s_s, None

    run_table(
        "Pipeline 4: distinct()",
        sizes,
        has_robin,
        has_sparkless,
        iterations,
        args.check_results,
        row4,
    )

    # ----- Pipeline 5: union two DataFrames then groupBy count (fixed size) -----
    n_union = min(UNION_SIZE, max(sizes)) if sizes else UNION_SIZE
    data_union_a = make_data(n_union)
    data_union_b = make_data(n_union)

    def row5(_n: int):
        r_m, r_s = (time_it(lambda: run_robin_union_then_groupby(data_union_a, data_union_b, cols), iterations) if has_robin else (None, None))
        s_m, s_s = (time_it(lambda: run_sparkless_union_then_groupby(data_union_a, data_union_b, cols), iterations) if has_sparkless else (None, None))
        return r_m, r_s, s_m, s_s, None

    run_table(
        f"Pipeline 5: union two dfs (n={n_union} each) -> groupBy(age).count()",
        [n_union * 2],  # total rows
        has_robin,
        has_sparkless,
        iterations,
        args.check_results,
        row5,
    )

    # ----- Pipeline 6: inner join (JOIN_SIZE x JOIN_SIZE) -----
    n_join = min(JOIN_SIZE, max(sizes)) if sizes else JOIN_SIZE
    data_join_a = [(i, i % 50, f"a_{i}") for i in range(n_join)]
    data_join_b = [(i, i * 2, f"b_{i}") for i in range(n_join)]

    def row6(_n: int):
        r_m, r_s = (time_it(lambda: run_robin_join(data_join_a, data_join_b), iterations) if has_robin else (None, None))
        s_m, s_s = (time_it(lambda: run_sparkless_join(data_join_a, data_join_b), iterations) if has_sparkless else (None, None))
        return r_m, r_s, s_m, s_s, None

    run_table(
        f"Pipeline 6: inner join (n={n_join} x {n_join})",
        [n_join],
        has_robin,
        has_sparkless,
        iterations,
        args.check_results,
        row6,
    )

    return 0


if __name__ == "__main__":
    sys.exit(main())
