#!/usr/bin/env python3
"""
Standalone benchmark for end-to-end WDM transform/inverse:

  time series -> t2w -> w2t -> reconstructed series

Compares local `wdm_wavelet` backend (`jax`) and the
pycwb ROOT-backed WDM implementation (if available).
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from time import perf_counter
from pathlib import Path
from typing import Dict, List

import numpy as np

# Ensure we import local workspace code, not an installed copy.
REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from wdm_wavelet.core import t2w as t2w_core
from wdm_wavelet.core import w2t as w2t_core
from wdm_wavelet.wdm import WDM


HAS_T2W_JAX = bool(getattr(t2w_core, "HAS_JAX", False))
HAS_T2W_NUMBA = bool(getattr(t2w_core, "HAS_NUMBA", False))
HAS_W2T_JAX = bool(getattr(w2t_core, "HAS_JAX", False))
HAS_W2T_NUMBA = bool(getattr(w2t_core, "HAS_NUMBA", False))


ROOT_CHILD_CODE = r"""
import argparse
import json
import os
from time import perf_counter

import numpy as np


def run(args):
    import ROOT
    from gwpy.timeseries import TimeSeries
    from pycwb.modules.cwb_conversions import convert_to_wavearray, convert_wavearray_to_nparray

    rng = np.random.default_rng(args.seed)
    signal = rng.standard_normal(args.length).astype(np.float64)
    signal_norm = float(np.linalg.norm(signal))

    def one_pass():
        # CWB path requested by user:
        #   wdm_cwb = ROOT.WDM(np.double)(M, K, beta_order, precision)
        #   tf_map_cwb = ROOT.WSeries(np.double)(wavearray, wdm_cwb)
        #   tf_map_cwb.Forward()
        #   tf_map_cwb.Inverse()
        hp = TimeSeries(signal, sample_rate=args.sample_rate, t0=0.0)
        wavearray = convert_to_wavearray(hp)

        wdm_cwb = ROOT.WDM(np.double)(args.M, args.K, args.beta_order, args.precision)
        tf_map_cwb = ROOT.WSeries(np.double)(wavearray, wdm_cwb)

        t0 = perf_counter()
        tf_map_cwb.Forward()
        t1 = perf_counter()
        tf_map_cwb.Inverse()
        t2 = perf_counter()

        rec = np.asarray(convert_wavearray_to_nparray(tf_map_cwb, clean=False), dtype=np.float64)
        rel_err = float(np.linalg.norm(rec - signal) / signal_norm)
        return (t1 - t0), (t2 - t1), (t2 - t0), rel_err

    for _ in range(max(0, int(args.warmup))):
        one_pass()

    t2w_times = []
    w2t_times = []
    total_times = []
    rel_err = None
    for _ in range(max(1, int(args.n_runs))):
        dt_t2w, dt_w2t, dt_total, rel_err = one_pass()
        t2w_times.append(dt_t2w)
        w2t_times.append(dt_w2t)
        total_times.append(dt_total)

    result = {
        "t2w_backend": "pycwb_root",
        "w2t_backend": "pycwb_root",
        "t2w_mean_s": float(np.mean(t2w_times)),
        "t2w_std_s": float(np.std(t2w_times)),
        "w2t_mean_s": float(np.mean(w2t_times)),
        "w2t_std_s": float(np.std(w2t_times)),
        "total_mean_s": float(np.mean(total_times)),
        "total_std_s": float(np.std(total_times)),
        "n_runs": int(args.n_runs),
        "reconstruction_rel_err": float(rel_err),
    }

    # Avoid ROOT/PyROOT teardown crashes by bypassing interpreter finalizers.
    print(json.dumps(result), flush=True)
    import sys
    sys.stdout.flush()
    os._exit(0)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--M", type=int, required=True)
    parser.add_argument("--K", type=int, required=True)
    parser.add_argument("--beta-order", dest="beta_order", type=int, required=True)
    parser.add_argument("--precision", type=int, required=True)
    parser.add_argument("--length", type=int, required=True)
    parser.add_argument("--sample-rate", dest="sample_rate", type=float, required=True)
    parser.add_argument("--seed", type=int, required=True)
    parser.add_argument("--MM", type=int, required=True)
    parser.add_argument("--n-runs", dest="n_runs", type=int, required=True)
    parser.add_argument("--warmup", type=int, required=True)
    run(parser.parse_args())
"""


def _available_python_backends() -> List[str]:
    backends = []
    if HAS_T2W_JAX and HAS_W2T_JAX:
        backends.append("jax")
    if HAS_T2W_NUMBA and HAS_W2T_NUMBA:
        backends.append("numba")
    return backends


def _parse_backends(backends_arg: str) -> List[str]:
    available = _available_python_backends()
    if backends_arg.strip().lower() == "all":
        return available

    requested = [b.strip().lower() for b in backends_arg.split(",") if b.strip()]
    valid = {"jax", "numba"}
    invalid = [b for b in requested if b not in valid]
    if invalid:
        raise ValueError(f"Unknown backends: {invalid}. Valid choices: jax,numba,all")

    unavailable = [b for b in requested if b not in available]
    if unavailable:
        raise ValueError(
            f"Requested backends not available in this environment: {unavailable}. "
            f"Available: {available}"
        )
    return requested


def _timeseries_values(ts):
    if hasattr(ts, "value"):
        return np.asarray(ts.value)
    return np.asarray(ts)


def benchmark_python_backends(
    signal: np.ndarray,
    sample_rate: float,
    M: int,
    K: int,
    beta_order: int,
    precision: int,
    MM: int,
    backends: List[str],
    n_runs: int,
    warmup: int,
) -> Dict[str, Dict[str, float]]:
    signal_norm = float(np.linalg.norm(signal))
    results: Dict[str, Dict[str, float]] = {}

    for backend in backends:
        wdm = WDM(M, K, beta_order, precision, backend=str(backend).lower())
        resolved_backend = wdm.backend

        for _ in range(max(0, int(warmup))):
            tf_warmup = wdm.t2w(signal, sample_rate=sample_rate, MM=MM)
            _ = wdm.w2t(tf_warmup)

        t2w_times = []
        w2t_times = []
        total_times = []
        rel_err = None

        for _ in range(max(1, int(n_runs))):
            start = perf_counter()
            tf_map = wdm.t2w(signal, sample_rate=sample_rate, MM=MM)
            mid = perf_counter()
            reconstructed_ts = wdm.w2t(tf_map)
            end = perf_counter()

            reconstructed = _timeseries_values(reconstructed_ts)
            rel_err = float(np.linalg.norm(reconstructed - signal) / signal_norm)

            t2w_times.append(mid - start)
            w2t_times.append(end - mid)
            total_times.append(end - start)

        results[str(backend)] = {
            "t2w_backend": resolved_backend,
            "w2t_backend": resolved_backend,
            "t2w_mean_s": float(np.mean(t2w_times)),
            "t2w_std_s": float(np.std(t2w_times)),
            "w2t_mean_s": float(np.mean(w2t_times)),
            "w2t_std_s": float(np.std(w2t_times)),
            "total_mean_s": float(np.mean(total_times)),
            "total_std_s": float(np.std(total_times)),
            "n_runs": int(n_runs),
            "reconstruction_rel_err": float(rel_err),
        }

    return results


def benchmark_pycwb_root_subprocess(
    M: int,
    K: int,
    beta_order: int,
    precision: int,
    length: int,
    sample_rate: float,
    seed: int,
    MM: int,
    n_runs: int,
    warmup: int,
) -> Dict[str, float]:
    cmd = [
        sys.executable,
        "-c",
        ROOT_CHILD_CODE,
        "--M",
        str(M),
        "--K",
        str(K),
        "--beta-order",
        str(beta_order),
        "--precision",
        str(precision),
        "--length",
        str(length),
        "--sample-rate",
        str(sample_rate),
        "--seed",
        str(seed),
        "--MM",
        str(MM),
        "--n-runs",
        str(n_runs),
        "--warmup",
        str(warmup),
    ]

    env = os.environ.copy()
    env.setdefault("MPLCONFIGDIR", "/tmp")

    proc = subprocess.run(cmd, capture_output=True, text=True, env=env, cwd=str(REPO_ROOT), check=False)

    parsed = None
    for line in reversed(proc.stdout.splitlines()):
        line = line.strip()
        if not line:
            continue
        try:
            parsed = json.loads(line)
            break
        except json.JSONDecodeError:
            continue

    if parsed is None:
        stderr_tail = "\n".join(proc.stderr.splitlines()[-12:])
        stdout_tail = "\n".join(proc.stdout.splitlines()[-12:])
        raise RuntimeError(
            "Failed to parse pycwb ROOT benchmark output.\n"
            f"returncode={proc.returncode}\n"
            f"stdout tail:\n{stdout_tail}\n"
            f"stderr tail:\n{stderr_tail}"
        )

    return parsed


def _format_ms(seconds: float) -> str:
    return f"{1000.0 * float(seconds):9.3f}"


def print_comparison_table(results: Dict[str, Dict[str, float]]):
    headers = [
        "backend",
        "t2w_ms",
        "w2t_ms",
        "total_ms",
        "total_std_ms",
        "rel_err",
    ]
    line = (
        f"{headers[0]:<14} {headers[1]:>10} {headers[2]:>10} "
        f"{headers[3]:>10} {headers[4]:>13} {headers[5]:>12}"
    )
    print(line)
    print("-" * len(line))

    for backend, s in results.items():
        print(
            f"{backend:<14} "
            f"{_format_ms(s['t2w_mean_s'])} "
            f"{_format_ms(s['w2t_mean_s'])} "
            f"{_format_ms(s['total_mean_s'])} "
            f"{_format_ms(s['total_std_s'])} "
            f"{s.get('reconstruction_rel_err', float('nan')):12.6g}"
        )


def build_arg_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Benchmark end-to-end t2w/w2t for jax and pycwb ROOT."
    )
    p.add_argument("--M", type=int, default=32, help="WDM max frequency layer index.")
    p.add_argument("--K", type=int, default=64, help="WDM K parameter.")
    p.add_argument("--beta-order", dest="beta_order", type=int, default=6, help="WDM beta order.")
    p.add_argument("--precision", type=int, default=10, help="WDM filter truncation precision.")
    p.add_argument("--length", type=int, default=4096, help="Signal length.")
    p.add_argument("--sample-rate", type=float, default=4096.0, help="Sample rate in Hz.")
    p.add_argument("--MM", type=int, default=-1, help="t2w MM mode/stride parameter.")
    p.add_argument("--seed", type=int, default=20260223, help="Random seed for input signal.")
    p.add_argument("--n-runs", type=int, default=5, help="Benchmark runs.")
    p.add_argument("--warmup", type=int, default=1, help="Warmup runs.")
    p.add_argument(
        "--backends",
        type=str,
        default="all",
        help="Comma-separated python backends (jax) or 'all'.",
    )
    p.add_argument(
        "--no-root",
        action="store_true",
        help="Skip pycwb ROOT benchmark.",
    )
    p.add_argument(
        "--json-out",
        type=str,
        default="",
        help="Optional path to write raw benchmark JSON.",
    )
    return p


def main() -> int:
    args = build_arg_parser().parse_args()
    backends = _parse_backends(args.backends)

    rng = np.random.default_rng(args.seed)
    signal = rng.standard_normal(args.length).astype(np.float64)

    print(
        f"Config: M={args.M} K={args.K} beta_order={args.beta_order} precision={args.precision} "
        f"length={args.length} sample_rate={args.sample_rate} MM={args.MM} "
        f"n_runs={args.n_runs} warmup={args.warmup} seed={args.seed}"
    )
    print(f"Python backends: {backends}")

    all_results: Dict[str, Dict[str, float]] = {}
    python_results = benchmark_python_backends(
        signal=signal,
        sample_rate=args.sample_rate,
        M=args.M,
        K=args.K,
        beta_order=args.beta_order,
        precision=args.precision,
        MM=args.MM,
        backends=backends,
        n_runs=args.n_runs,
        warmup=args.warmup,
    )
    all_results.update(python_results)

    if not args.no_root:
        try:
            root_result = benchmark_pycwb_root_subprocess(
                M=args.M,
                K=args.K,
                beta_order=args.beta_order,
                precision=args.precision,
                length=args.length,
                sample_rate=args.sample_rate,
                seed=args.seed,
                MM=args.MM,
                n_runs=args.n_runs,
                warmup=args.warmup,
            )
            all_results["pycwb_root"] = root_result
        except Exception as exc:
            print(f"pycwb_root benchmark failed: {exc}")

    print()
    print_comparison_table(all_results)

    if args.json_out:
        out_path = Path(args.json_out).expanduser()
        out_path.parent.mkdir(parents=True, exist_ok=True)
        with out_path.open("w", encoding="utf-8") as f:
            json.dump(all_results, f, indent=2, sort_keys=True)
        print(f"\nSaved JSON: {out_path}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
