# type: ignore
"""
Benchmark: PyArrow-based to_dataframe() vs. legacy CSV-roundtrip to_dataframe().

Fetches data once from the API, caches it as a list of dicts, then repeatedly
converts that cached data to a DataFrame using both methods so that
network I/O is excluded from the timing.
"""
import os
import sys
import time
import statistics
from copy import deepcopy
from io import BytesIO
from typing import Any, Dict, List, Optional

import numpy as np
import pandas as pd
import polars as pl
import pyarrow as pa
import pyarrow.compute as pc

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from coinmetrics.api_client import CoinMetricsClient
from coinmetrics._data_collection import DataCollection


ITERATIONS = 5


# ---------------------------------------------------------------------------
# Legacy method (CSV round-trip, as it existed on master)
# ---------------------------------------------------------------------------
def legacy_to_dataframe(
    rows: List[Dict[str, Any]],
    dtype_mapper: Optional[Dict[str, Any]] = None,
) -> pd.DataFrame:
    """Reproduce the old optimize_dtypes=True CSV round-trip path."""
    if not rows:
        return pd.DataFrame()

    columns = list(rows[0].keys())
    lines = [",".join(columns) + "\n"]
    for row in rows:
        lines.append(
            ",".join(f'"{row.get(c)}"' for c in columns) + "\n"
        )
    csv_bytes = "".join(lines).encode()
    buf = BytesIO(csv_bytes)

    datetime_cols = [c for c in columns if c.endswith("_time") or c == "time"]
    df = pd.read_csv(buf, parse_dates=datetime_cols, dtype=dtype_mapper)
    if dtype_mapper is None:
        df = df.convert_dtypes()
    if df.dtypes.get("coin_metrics_id") == np.dtype("object"):
        df["coin_metrics_id"] = df["coin_metrics_id"].astype(np.float128)
    return df


def legacy_to_polars(rows: List[Dict[str, Any]]) -> pl.DataFrame:
    """Reproduce the old polars path: CSV round-trip → pl.read_csv."""
    if not rows:
        return pl.DataFrame()

    columns = list(rows[0].keys())
    lines = [",".join(columns) + "\n"]
    for row in rows:
        lines.append(
            ",".join(f'"{row.get(c)}"' for c in columns) + "\n"
        )
    csv_bytes = "".join(lines).encode()
    return pl.read_csv(BytesIO(csv_bytes), infer_schema_length=10000)


# ---------------------------------------------------------------------------
# New method (PyArrow from_pylist + per-column cast)
# ---------------------------------------------------------------------------
def pyarrow_to_dataframe(
    rows: List[Dict[str, Any]],
    dc: DataCollection,
    dtype_mapper: Optional[Dict[str, Any]] = None,
    decimal_as_string: bool = False,
) -> pd.DataFrame:
    """Reproduce the current optimized PyArrow path on cached rows."""
    if not rows:
        pa_schema = dc._item_schema.get_arrow_schema()
        empty_table = pa.table(
            {f.name: pa.array([], type=f.type) for f in pa_schema}
        )
        empty_table = DataCollection._prepare_table_for_pandas(
            empty_table, decimal_as_string=decimal_as_string
        )
        return empty_table.to_pandas(types_mapper=DataCollection._pandas_type_mapper)

    pa_schema = dc._item_schema.get_arrow_schema()

    batch_size = 10_000
    batches: list[pa.RecordBatch] = []
    batch_rows: list[dict] = []

    for row in rows:
        batch_rows.append(row)
        if len(batch_rows) >= batch_size:
            batches.append(DataCollection._flush_batch_rows(
                batch_rows, pa_schema, skip_nested_cast=True
            ))
            batch_rows.clear()

    if batch_rows:
        batches.append(DataCollection._flush_batch_rows(
            batch_rows, pa_schema, skip_nested_cast=True
        ))

    tables = [pa.Table.from_batches([b]) for b in batches]
    table = pa.concat_tables(tables, promote_options="permissive")
    table = DataCollection._prepare_table_for_pandas(
        table, decimal_as_string=decimal_as_string
    )
    df = table.to_pandas(types_mapper=DataCollection._pandas_type_mapper)
    if dtype_mapper:
        applicable = {k: v for k, v in dtype_mapper.items() if k in df.columns}
        if applicable:
            df = df.astype(applicable)
    return df


def pyarrow_to_polars(
    rows: List[Dict[str, Any]],
    dc: DataCollection,
) -> pl.DataFrame:
    """Reproduce the current PyArrow → polars path on cached rows."""
    if not rows:
        pa_schema = dc._item_schema.get_arrow_schema()
        empty_table = pa.table(
            {f.name: pa.array([], type=f.type) for f in pa_schema}
        )
        return pl.from_arrow(empty_table)

    pa_schema = dc._item_schema.get_arrow_schema()

    batch_size = 10_000
    batches: list[pa.RecordBatch] = []
    batch_rows: list[dict] = []

    for row in rows:
        batch_rows.append(row)
        if len(batch_rows) >= batch_size:
            batches.append(DataCollection._flush_batch_rows(
                batch_rows, pa_schema, skip_nested_cast=False
            ))
            batch_rows.clear()

    if batch_rows:
        batches.append(DataCollection._flush_batch_rows(
            batch_rows, pa_schema, skip_nested_cast=False
        ))

    table = pa.Table.from_batches(batches, schema=pa_schema)
    return pl.from_arrow(table)


# ---------------------------------------------------------------------------
# Runner
# ---------------------------------------------------------------------------
def benchmark(label: str, func, iterations: int = ITERATIONS):
    """Run *func* multiple times and report timing statistics."""
    times = []
    result = None
    for i in range(iterations):
        t0 = time.perf_counter()
        result = func()
        elapsed = time.perf_counter() - t0
        times.append(elapsed)
        print(f"  {label} run {i + 1}/{iterations}: {elapsed:.4f}s")

    mean = statistics.mean(times)
    median = statistics.median(times)
    stdev = statistics.stdev(times) if len(times) > 1 else 0.0
    best = min(times)
    worst = max(times)
    print(f"  {label} -- mean: {mean:.4f}s  median: {median:.4f}s  "
          f"stdev: {stdev:.4f}s  best: {best:.4f}s  worst: {worst:.4f}s")
    print(f"  {label} -- result shape: {result.shape}")
    print()
    return mean, result


def run_benchmark(
    name: str, dc: DataCollection, dataframe_type: str = "pandas"
):
    print(f"\n{'='*72}")
    print(f"  Fetching data for: {name}  [{dataframe_type}]")
    print(f"{'='*72}")

    t0 = time.perf_counter()
    rows = deepcopy(dc).to_list()
    fetch_time = time.perf_counter() - t0
    print(f"  Fetched {len(rows):,} rows in {fetch_time:.4f}s\n")

    if dataframe_type == "polars":
        print("--- Legacy (CSV round-trip → polars) ---")
        legacy_mean, _ = benchmark(
            "legacy", lambda: legacy_to_polars(rows),
        )
        print("--- New (PyArrow → polars) ---")
        arrow_mean, _ = benchmark(
            "pyarrow", lambda: pyarrow_to_polars(rows, dc),
        )
    else:
        print("--- Legacy (CSV round-trip) ---")
        legacy_mean, _ = benchmark(
            "legacy", lambda: legacy_to_dataframe(rows),
        )
        print("--- New (PyArrow from_pylist) ---")
        arrow_mean, _ = benchmark(
            "pyarrow", lambda: pyarrow_to_dataframe(rows, dc),
        )

    speedup = legacy_mean / arrow_mean if arrow_mean > 0 else float("inf")
    print(f"{'='*72}")
    print(f"  {name}  [{dataframe_type}]")
    print(f"    Legacy:  {legacy_mean:.4f}s")
    print(f"    PyArrow: {arrow_mean:.4f}s")
    print(f"    Speedup: {speedup:.2f}x")
    print(f"{'='*72}\n")

    return {
        "name": name,
        "type": dataframe_type,
        "rows": len(rows),
        "legacy_mean": legacy_mean,
        "arrow_mean": arrow_mean,
        "speedup": speedup,
    }


if __name__ == "__main__":
    client = CoinMetricsClient(os.environ.get("CM_API_KEY"))

    benchmarks = [
        (
            "reference_data_markets",
            client.reference_data_markets(),
            "pandas",
        ),
        (
            "get_asset_metrics(5 assets, ReferenceRateUSD, 1s, lim=100k/asset)",
            client.get_asset_metrics(
                assets=["btc", "eth", "xrp", "doge", "sol"],
                metrics="ReferenceRateUSD",
                frequency="1s", limit_per_asset=100000,
            ),
            "pandas",
        ),
        (
            "get_market_orderbooks(coinbase-btc-usd-spot, raw, lim=10k)",
            client.get_market_orderbooks(
                "coinbase-btc-usd-spot", granularity="raw",
                start_time="2025-10-10", limit_per_market=10000,
            ),
            "pandas",
        ),
        (
            "get_market_trades(coinbase+binance-btc, lim=100k/mkt)",
            client.get_market_trades(
                "coinbase-btc-usd-spot,binance-BTCUSDT-future",
                limit_per_market=100_000,
            ),
            "pandas",
        ),
        (
            "get_list_of_blocks_v2(btc, 800k-900k blocks)",
            client.get_list_of_blocks_v2("btc", start_height=800_000, end_height=900_000, page_size=10000),
            "polars",
        ),
        (
            "get_list_of_balance_updates_v2(btc, fees, 900k-900.5k)",
            client.get_list_of_balance_updates_v2(
                "btc", include_sub_accounts=True, accounts='FEES',
                start_height=900_000, end_height=900_500, page_size=10000
            ),
            "polars",
        ),
    ]

    results = []
    for name, dc, dftype in benchmarks:
        results.append(run_benchmark(name, dc, dataframe_type=dftype))

    print("\n" + "=" * 72)
    print("  SUMMARY")
    print("=" * 72)
    for r in results:
        label = f"{r['name']} [{r['type']}]"
        print(f"  {label:72s}  rows={r['rows']:>10,}  "
              f"legacy={r['legacy_mean']:.4f}s  arrow={r['arrow_mean']:.4f}s  "
              f"speedup={r['speedup']:.2f}x")
    print("=" * 72)
