"""
Benchmark: Run 100 mutations through the full mutation pipeline and report
per-mutation timing for each pipeline stage.

Usage:
    python tests/benchmark_100.py
    python tests/benchmark_100.py --n 50 --with-borzoi --with-mirna
    python tests/benchmark_100.py --mutations-file my_mutations.txt

Output: TSV with per-mutation timing breakdown + summary statistics.
"""

import argparse
import csv
import sys
import time

# 100 diverse mutations across cancer genes, variant types, and genomic regions
DEFAULT_MUTATIONS = [
    "KRAS:12:25227343:G:T", "BRAF:7:140753336:A:T", "TP53:17:7673802:C:T",
    "BRCA1:17:43094464:G:A", "BRCA2:13:32340300:G:T", "EGFR:7:55259515:T:G",
    "MET:7:116412043:A:G", "PIK3CA:3:179234297:A:G", "PTEN:10:87933147:C:T",
    "APC:5:112175770:G:A", "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",
    "RB1:13:48367512:C:T", "VHL:3:10183842:G:A", "CDH1:16:68835639:G:A",
    "SMAD4:18:51065544:G:A", "CTNNB1:3:41266098:C:T", "MYC:8:128750986:G:A",
    "CDKN2A:9:21968228:C:T", "IDH1:2:208248389:G:A", "NRAS:1:114713909:G:A",
    "HRAS:11:534288:G:A", "ALK:2:29220826:G:A", "RET:10:43613908:G:T",
    "FGFR3:4:1801841:G:A", "KIT:4:55593464:A:T", "PDGFRA:4:54285925:A:G",
    "KRAS:12:25227342:C:T", "TP53:17:7674220:C:T", "EGFR:7:55259516:G:A",
    "BRAF:7:140753337:T:C", "PIK3CA:3:179234298:G:A", "TP53:17:7675100:G:A",
    "BRCA1:17:43095000:C:T", "RB1:13:48370000:G:A", "APC:5:112176000:T:C",
    "VHL:3:10184000:C:T", "NPM1:5:171410540:C:T", "FLT3:13:28034105:G:A",
    "DNMT3A:2:25234374:G:A", "TET2:4:105236048:G:A", "JAK2:9:5073770:G:T",
    "ABL1:9:130872896:T:C", "ERBB2:17:39724745:G:A", "NOTCH1:9:136504900:G:A",
    "FBXW7:4:152326166:G:A", "SF3B1:2:197402110:G:A", "SRSF2:17:76736877:C:T",
    "U2AF1:21:43092956:T:C", "ZRSR2:X:15841214:G:A", "EZH2:7:148506440:G:A",
    "KMT2A:11:118377655:G:A", "CREBBP:16:3775000:G:A", "EP300:22:41151000:G:A",
    "ARID1A:1:26779000:G:A", "SMARCA4:19:10964000:G:A", "STAG2:X:123185000:G:A",
    "RAD21:8:116845000:G:A", "SMC1A:X:53432000:G:A", "SMC3:10:111686000:G:A",
    "ASXL1:20:32434000:G:A", "BCOR:X:39920000:G:A", "BCORL1:X:129147000:G:A",
    "PHF6:X:133547000:G:A", "RUNX1:21:34792000:G:A", "GATA2:3:128202000:G:A",
    "CEBPA:19:33301000:G:A", "WT1:11:32417000:G:A", "NF1:17:31259000:G:A",
    "NF2:22:29999000:G:A", "TSC1:9:134772000:G:A", "TSC2:16:2097000:G:A",
    "STK11:19:1170000:G:A", "LZTR1:22:20455000:G:A", "MAP2K1:15:66386000:G:A",
    "MAP2K2:19:4090000:G:A", "ARAF:X:47422000:G:A", "RAF1:3:12626000:G:A",
    "PTPN11:12:112450000:G:A", "SOS1:2:39208000:G:A", "CBL:11:119149000:G:A",
    "CBLB:3:105635000:G:A", "PPM1D:17:60663000:G:A", "CHEK2:22:28725000:G:A",
    "ATM:11:108098000:G:A", "ATR:3:142168000:G:A", "PALB2:16:23603000:G:A",
    "RAD51C:17:58697000:G:A", "MLH1:3:37034000:G:A", "MSH2:2:47630000:G:A",
    "MSH6:2:47783000:G:A", "PMS2:7:5979000:G:A", "POLE:12:132624000:G:A",
    "POLD1:19:50384000:G:A", "CDK4:12:57750000:G:A", "CCND1:11:69641000:G:A",
    "MDM2:12:68808000:G:A", "CDK6:7:92463000:G:A", "CDKN1B:12:12716000:G:A",
    "FGFR1:8:38411000:G:A", "FGFR2:10:121551000:G:A",
]


def run_benchmark(mutations, run_borzoi=False, run_mirna=False):
    from geney import mutation_pipeline

    results = []
    n = len(mutations)

    print(f"Benchmarking {n} mutations")
    print(f"  Borzoi: {'ON' if run_borzoi else 'OFF'}")
    print(f"  miRNA:  {'ON' if run_mirna else 'OFF'}")
    print(f"{'─'*80}")

    for i, mut_id in enumerate(mutations):
        gene = mut_id.split(":")[0]
        t0 = time.perf_counter()

        try:
            result = mutation_pipeline(
                mut_id,
                run_transcription=run_borzoi,
                run_mirna=run_mirna,
            )
            elapsed = time.perf_counter() - t0
            t = result.timing

            row = {
                "mutation": mut_id,
                "gene": gene,
                "total": round(elapsed, 3),
                "prepare": round(t.get("prepare", 0), 3),
                "spliceai": round(t.get("spliceai", 0), 3),
                "protein": round(t.get("protein", 0), 3),
                "tis": round(t.get("tis", 0), 3),
                "folding": round(t.get("folding", 0), 3),
                "borzoi": round(t.get("borzoi", 0), 3),
                "mirna": round(t.get("mirna", 0), 3),
                "score": result.overall_score,
                "isoforms": len(result.splicing.isoforms),
                "protein_type": result.protein.mutation_type if result.protein else "",
                "errors": ",".join(result.errors.keys()),
            }
        except Exception as e:
            elapsed = time.perf_counter() - t0
            row = {"mutation": mut_id, "gene": gene, "total": round(elapsed, 3),
                   "errors": str(e)[:60]}

        results.append(row)
        score = row.get("score", "?")
        print(f"  [{i+1:3d}/{n}] {gene:12s} {elapsed:6.2f}s  score={score}")

    return results


def summarize(results):
    import statistics
    ok = [r for r in results if "score" in r]
    if not ok:
        print("\nNo successful results.")
        return

    totals = [r["total"] for r in ok]
    print(f"\n{'═'*60}")
    print(f"  {len(ok)}/{len(results)} succeeded")
    print(f"  Total:  {sum(totals):.1f}s")
    print(f"  Mean:   {statistics.mean(totals):.2f}s")
    print(f"  Median: {statistics.median(totals):.2f}s")
    print(f"  Min:    {min(totals):.2f}s  Max: {max(totals):.2f}s")
    if len(totals) > 1:
        print(f"  Stdev:  {statistics.stdev(totals):.2f}s")

    stages = ["prepare", "spliceai", "protein", "tis", "folding", "borzoi", "mirna"]
    print(f"\n  Stage breakdown (mean seconds):")
    for s in stages:
        vals = [r.get(s, 0) for r in ok]
        if sum(vals) > 0:
            print(f"    {s:12s}: {statistics.mean(vals):.3f}s")
    print(f"{'═'*60}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--n", type=int, default=100)
    parser.add_argument("--with-borzoi", action="store_true")
    parser.add_argument("--with-mirna", action="store_true")
    parser.add_argument("--output", default="benchmark_100_results.tsv")
    args = parser.parse_args()

    muts = DEFAULT_MUTATIONS[:args.n]
    results = run_benchmark(muts, args.with_borzoi, args.with_mirna)
    summarize(results)

    keys = list(results[0].keys())
    with open(args.output, "w", newline="") as f:
        w = csv.DictWriter(f, keys, delimiter="\t", extrasaction="ignore")
        w.writeheader()
        w.writerows(results)
    print(f"\nSaved to {args.output}")
