# tests/profile_pipeline.py — Profile where time is spent in the full pipeline.
import time
import cProfile
import pstats
from io import StringIO
from geney.pipelines import mutation_pipeline

MUT = "BRCA1:17:43047640:T:C"

# Warm up models first
print("Warming up...")
mutation_pipeline(MUT, run_splicing=True, run_tis=True, run_folding=True, run_transcription=False)

print(f"\nProfiling full pipeline on {MUT} (no Borzoi)...\n")

pr = cProfile.Profile()
pr.enable()
result = mutation_pipeline(MUT, run_splicing=True, run_tis=True, run_folding=True, run_transcription=False)
pr.disable()

s = StringIO()
ps = pstats.Stats(pr, stream=s)
ps.sort_stats('cumulative')
ps.print_stats(40)
print(s.getvalue())

# Manual timing of each phase
print("\n--- Manual phase timing ---")

from geney.variants import MutationalEvent
from geney.utils import select_transcript, _apply_mutation_safe
from seqmat import Gene

m = MutationalEvent(MUT)

t0 = time.perf_counter()
gene = Gene.from_file(m.gene, organism="hg38")
print(f"Gene load:          {time.perf_counter()-t0:.4f}s")

t0 = time.perf_counter()
selected = select_transcript(gene, m.central_position)
ref_transcript = selected.generate_pre_mrna().generate_mature_mrna().generate_protein()
print(f"Ref transcript:     {time.perf_counter()-t0:.4f}s")

t0 = time.perf_counter()
var_transcript = selected.clone().generate_pre_mrna()
for pos, ref, alt in m.mutation_args():
    _apply_mutation_safe(var_transcript.pre_mrna, pos, ref, alt)
var_transcript.generate_mature_mrna().generate_protein()
print(f"Var transcript:     {time.perf_counter()-t0:.4f}s")

# Splicing
from geney.splicing import TranscriptLibrary, SpliceSimulator
from geney.oncosplice import Oncosplice

t0 = time.perf_counter()
tl = TranscriptLibrary(ref_transcript, m, individual_mutations=False)
print(f"TranscriptLibrary:  {time.perf_counter()-t0:.4f}s")

t0 = time.perf_counter()
tl.predict_splicing(m.central_position, engine="spliceai", inplace=True)
print(f"predict_splicing:   {time.perf_counter()-t0:.4f}s")

t0 = time.perf_counter()
splicing_results = tl.get_event_columns("event")
ss = SpliceSimulator(splicing_results, tl.event, feature="event", max_distance=100_000_000)
print(f"SpliceSimulator:    {time.perf_counter()-t0:.4f}s")

t0 = time.perf_counter()
for variant_transcript, isoform_metadata in ss.get_viable_transcripts(metadata=True):
    onco = Oncosplice(ref_transcript.protein, variant_transcript.protein, None)
    _ = onco.get_analysis_series()
print(f"Oncosplice loop:    {time.perf_counter()-t0:.4f}s")

# TIS
t0 = time.perf_counter()
from geney.tis import compute_tis_shift_metrics
from geney.tis.helpers import _get_mature_mrna_seq, _apply_mutation_to_mrna
ref_mrna_seq = _get_mature_mrna_seq(ref_transcript)
mrna = ref_transcript.mature_mrna
var_mrna_seq = _apply_mutation_to_mrna(ref_mrna_seq, mrna.index, m.mutation_args())
tis_result = compute_tis_shift_metrics(ref_mrna=ref_mrna_seq, var_mrna=var_mrna_seq)
print(f"TIS full:           {time.perf_counter()-t0:.4f}s")

# Folding
t0 = time.perf_counter()
from geney.folding.delta_mfe import compute_folding_profile
from geney.folding.reference import complement
ref_seq = mrna.seq.upper()
idx_list = mrna.index.tolist()
pos_to_mrna = {gpos: i for i, gpos in enumerate(idx_list)}
for pos, ref_allele, alt_allele in m.mutation_args():
    mrna_pos = pos_to_mrna.get(pos)
    if mrna_pos is not None:
        actual_ref = ref_seq[mrna_pos]
        coding_alt = alt_allele.upper()
        if actual_ref != ref_allele.upper():
            coding_alt = complement(alt_allele.upper())
        var_seq = ref_seq[:mrna_pos] + coding_alt + ref_seq[mrna_pos + 1:]
        profile = compute_folding_profile(ref_seq, var_seq, mrna_pos, 39)
        break
print(f"Folding full:       {time.perf_counter()-t0:.4f}s")
