# tests/benchmark_screen.py — Benchmark screen vs full pipeline.
"""
Compare wall-clock time: screen_mutation vs mutation_pipeline
for the same set of mutations, using real ClinVar data.
"""
import time
import duckdb
from geney.screen import screen_mutation, screen_mutations
from geney.pipelines import mutation_pipeline

CLINVAR_DB = "/Users/nicolaslynn/Downloads/clinvar/clinvar_tiles.duckdb"


def load_clinvar_sample(n=10, seed=42):
    """Grab n random mutations from ClinVar, filtering to common genes."""
    con = duckdb.connect(CLINVAR_DB, read_only=True)
    rows = con.execute(f"""
        SELECT mutation_id FROM clinvar
        WHERE LENGTH(ref) = 1 AND LENGTH(alt) = 1
        AND gene IN ('BRCA1','BRCA2','TP53','KRAS','EGFR','BRAF','PIK3CA','APC','PTEN','MLH1')
        ORDER BY hash(mutation_id || '{seed}')
        LIMIT {n}
    """).fetchall()
    con.close()
    return [r[0] for r in rows]


def bench_full_pipeline(muts, kwargs):
    results = []
    for mid in muts:
        try:
            r = mutation_pipeline(mid, **kwargs)
            results.append(r)
        except Exception as e:
            print(f"  full pipeline error on {mid}: {e}")
    return results


def bench_screen_individual(muts, kwargs):
    results = []
    for mid in muts:
        r = screen_mutation(mid, **kwargs)
        results.append(r)
    return results


def bench_screen_batch(muts, kwargs):
    return list(screen_mutations(muts, batch_size=len(muts), **kwargs))


def run_benchmark(muts, label, kwargs):
    n = len(muts)
    print(f"\n{'='*60}")
    print(f"{label} — {n} mutations")
    print(f"{'='*60}")

    # Warm up
    print("Warming up models...")
    screen_mutation(muts[0], **kwargs)
    mutation_pipeline(muts[0], **kwargs)
    print()

    # Full pipeline
    t0 = time.perf_counter()
    full_results = bench_full_pipeline(muts, kwargs)
    t_full = time.perf_counter() - t0
    print(f"Full pipeline:       {t_full:.3f}s  ({t_full/n:.3f}s/mut)")

    # Screen individual
    t0 = time.perf_counter()
    screen_results = bench_screen_individual(muts, kwargs)
    t_screen = time.perf_counter() - t0
    print(f"Screen (individual): {t_screen:.3f}s  ({t_screen/n:.3f}s/mut)")

    # Screen batched
    t0 = time.perf_counter()
    batch_results = bench_screen_batch(muts, kwargs)
    t_batch = time.perf_counter() - t0
    print(f"Screen (batched):    {t_batch:.3f}s  ({t_batch/n:.3f}s/mut)")

    print(f"\nSpeedup (individual screen vs full): {t_full/t_screen:.1f}x")
    print(f"Speedup (batched screen vs full):    {t_full/t_batch:.1f}x")

    print("\n--- Screen results ---")
    for r in batch_results:
        print(f"  {r.mut_id}: spl={r.splicing_delta}, tis={r.tis_delta}, "
              f"fold={r.folding_delta}, tx={r.transcription_delta}, "
              f"max={r.max_delta:.4f}" + (f" ERR={r.error}" if r.error else ""))


if __name__ == "__main__":
    muts = load_clinvar_sample(n=10)
    print(f"Sampled {len(muts)} ClinVar mutations:")
    for m in muts:
        print(f"  {m}")

    # Round 1: without Borzoi
    run_benchmark(muts, "WITHOUT Borzoi", dict(
        run_splicing=True, run_tis=True, run_folding=True, run_transcription=False,
    ))

    # Round 2: with Borzoi
    run_benchmark(muts, "WITH Borzoi", dict(
        run_splicing=True, run_tis=True, run_folding=True, run_transcription=True,
    ))
