# tests/benchmark_pipeline.py — Correctness + throughput benchmark for pipelines.
#
# Compares serial mutation_pipeline vs batch mutation_pipeline_batch:
#   1. Correctness: checks that both produce equivalent results.
#   2. Throughput: mutations/second for serial vs batch (no Borzoi).
#
# Run:  python tests/benchmark_pipeline.py

import time
import math
import numpy as np

from geney import mutation_pipeline, mutation_pipeline_batch

MUTATIONS = [
    "ECH1:19:38831664:A:C",
    "WBP4:13:41076038:G:C",
    "PRKACB:1:84191312:G:C",
    "TCF25:16:89895209:T:G",
    "TNPO3:7:128990171:G:A",
    "KRAS:12:25227343:G:T",
]

PIPELINE_KWARGS = dict(
    run_splicing=True,
    run_tis=True,
    run_folding=True,
    run_transcription=False,
)


def _compare_splicing(s1, s2, label):
    """Compare two SplicingResult objects."""
    ok = True
    if s1.empty != s2.empty:
        print(f"  FAIL [{label}] splicing empty mismatch: serial={s1.empty}, batch={s2.empty}")
        return False
    if s1.empty:
        return True
    if not math.isclose(s1.score, s2.score, abs_tol=1e-6):
        print(f"  FAIL [{label}] splicing score: serial={s1.score:.6f}, batch={s2.score:.6f}")
        ok = False
    if not math.isclose(s1.max_oncosplice_score, s2.max_oncosplice_score, abs_tol=1e-6):
        print(f"  FAIL [{label}] max_oncosplice: serial={s1.max_oncosplice_score:.6f}, batch={s2.max_oncosplice_score:.6f}")
        ok = False
    if len(s1.isoforms) != len(s2.isoforms):
        print(f"  FAIL [{label}] isoform count: serial={len(s1.isoforms)}, batch={len(s2.isoforms)}")
        ok = False
    return ok


def _compare_tis(t1, t2, label):
    """Compare two TISResult objects."""
    if t1 is None and t2 is None:
        return True
    if (t1 is None) != (t2 is None):
        print(f"  FAIL [{label}] tis None mismatch")
        return False
    ok = True
    if t1.status != t2.status:
        print(f"  FAIL [{label}] tis status: serial={t1.status}, batch={t2.status}")
        ok = False
    if not math.isclose(t1.score or 0, t2.score or 0, abs_tol=1e-6):
        print(f"  FAIL [{label}] tis score: serial={t1.score}, batch={t2.score}")
        ok = False
    return ok


def _compare_folding(f1, f2, label):
    """Compare two FoldingResult objects."""
    if f1 is None and f2 is None:
        return True
    if (f1 is None) != (f2 is None):
        print(f"  FAIL [{label}] folding None mismatch")
        return False
    ok = True
    if f1.status != f2.status:
        print(f"  FAIL [{label}] folding status: serial={f1.status}, batch={f2.status}")
        ok = False
    if not math.isclose(f1.score or 0, f2.score or 0, abs_tol=1e-6):
        print(f"  FAIL [{label}] folding score: serial={f1.score}, batch={f2.score}")
        ok = False
    if not math.isclose(f1.delta_mfe or 0, f2.delta_mfe or 0, abs_tol=1e-6):
        print(f"  FAIL [{label}] folding delta_mfe: serial={f1.delta_mfe}, batch={f2.delta_mfe}")
        ok = False
    return ok


def check_correctness(serial_results, batch_results):
    """Compare serial and batch results for equivalence."""
    assert len(serial_results) == len(batch_results)
    all_ok = True
    for i, (s, b) in enumerate(zip(serial_results, batch_results)):
        label = MUTATIONS[i].split(":")[0]
        ok = True
        ok &= _compare_splicing(s.splicing, b.splicing, label)
        ok &= _compare_tis(s.tis, b.tis, label)
        ok &= _compare_folding(s.folding, b.folding, label)
        if ok:
            print(f"  OK   [{label}]")
        all_ok &= ok
    return all_ok


def run_benchmark(n_repeats=3):
    """Time serial vs batch throughput."""
    n_muts = len(MUTATIONS)

    # -- Warmup (gene cache, model loads) --
    print("Warming up (first run loads models + gene cache)...")
    mutation_pipeline(MUTATIONS[0], **PIPELINE_KWARGS)
    mutation_pipeline_batch(MUTATIONS[:1], **PIPELINE_KWARGS)
    print()

    # -- Serial --
    serial_times = []
    serial_results = None
    for rep in range(n_repeats):
        t0 = time.perf_counter()
        results = [mutation_pipeline(m, **PIPELINE_KWARGS) for m in MUTATIONS]
        elapsed = time.perf_counter() - t0
        serial_times.append(elapsed)
        if serial_results is None:
            serial_results = results
        print(f"  Serial  run {rep+1}: {elapsed:.2f}s  ({n_muts/elapsed:.2f} mut/s)")

    # -- Batch --
    batch_times = []
    batch_results = None
    for rep in range(n_repeats):
        t0 = time.perf_counter()
        results = mutation_pipeline_batch(MUTATIONS, **PIPELINE_KWARGS)
        elapsed = time.perf_counter() - t0
        batch_times.append(elapsed)
        if batch_results is None:
            batch_results = results
        print(f"  Batch   run {rep+1}: {elapsed:.2f}s  ({n_muts/elapsed:.2f} mut/s)")

    # -- Summary --
    serial_avg = np.mean(serial_times)
    batch_avg = np.mean(batch_times)

    print(f"\n{'='*50}")
    print(f"Serial avg: {serial_avg:.2f}s  ({n_muts/serial_avg:.2f} mut/s)")
    print(f"Batch  avg: {batch_avg:.2f}s  ({n_muts/batch_avg:.2f} mut/s)")
    print(f"Speedup:    {serial_avg/batch_avg:.2f}x")

    # -- Timing breakdown (from last batch run) --
    if batch_results and batch_results[0].timing:
        print(f"\nBatch timing breakdown:")
        for phase, secs in batch_results[0].timing.items():
            print(f"  {phase:>12s}: {secs:.4f}s")

    # -- Correctness --
    print(f"\nCorrectness check (serial vs batch):")
    ok = check_correctness(serial_results, batch_results)
    print(f"\n{'PASS' if ok else 'FAIL'}")

    # -- Scores --
    print(f"\nScores (serial):")
    print(f"  {'Mutation':>8s}  {'splicing':>8s}  {'tis':>8s}  {'folding':>8s}  {'transcr':>8s}  {'overall':>8s}")
    for i, r in enumerate(serial_results):
        label = MUTATIONS[i].split(":")[0]
        sc = r.scores
        print(f"  {label:>8s}  {sc['splicing'].score:>8.1f}  {sc['tis'].score:>8.1f}  {sc['folding'].score:>8.1f}  {sc['transcription'].score:>8.1f}  {r.overall_score:>8.1f}")

    # -- Score consistency serial vs batch --
    scores_ok = True
    for i, (s, b) in enumerate(zip(serial_results, batch_results)):
        label = MUTATIONS[i].split(":")[0]
        if s.scores != b.scores:
            print(f"  FAIL [{label}] scores mismatch: serial={s.scores}, batch={b.scores}")
            scores_ok = False
    print(f"\nScore consistency: {'PASS' if scores_ok else 'FAIL'}")

    return ok and scores_ok


if __name__ == "__main__":
    run_benchmark()
