# geney/pipelines/_batch.py — Batch-optimized mutation pipeline.
from __future__ import annotations

import logging
import time
from concurrent.futures import ThreadPoolExecutor

import numpy as np

from ..results import (
    SplicingResult,
    TISResult,
    FoldingResult,
    TranscriptionResult,
    MutationResult,
)
from ..splicing import SpliceSimulator, TranscriptLibrary

from ._common import (
    SPLICING_CONTEXT_BP,
    SPLICE_MAX_DISTANCE,
    _load_gene_with_timeout,
    _prepare_mutation,
    _build_splice_site_index,
    _splicing_base_report,
    _oncosplice_dataframe,
    _prepare_tis,
    _tis_completed_dict,
    _run_folding_single,
    _transcription_result_from_dict,
)

logger = logging.getLogger(__name__)


def mutation_pipeline_batch(
    mut_ids: list[str],
    transcript_id: str | None = None,
    organism: str = "hg38",
    splicing_engine: str = "ensemble",
    context_bp: int = 1000,
    run_splicing: bool = True,
    run_tis: bool = True,
    run_folding: bool = True,
    run_transcription: bool = True,
    tis_proximity_window: int = 500,
    tis_probability_threshold: float = 0.1,
    folding_window_size: int = 39,
    gene_load_timeout: float | None = None,
    max_gpu_batch: int | None = None,
    n_folding_workers: int = 4,
    n_io_workers: int = 8,
    max_isoforms: int | None = None,
    min_total_prob: float = 0.5,
    min_delta_magnitude: float = 0.2,
) -> list[MutationResult]:
    """Batch-optimized mutation pipeline.

    Returns identical results to calling :func:`mutation_pipeline` N times,
    but internally batches GPU inference across mutations for higher throughput.

    Execution phases (CPU and GPU work overlapped where possible):
    1. PREPARE — Parse mutations, deduplicate gene loads, build transcripts.
    2. Launch FOLDING (CPU) + BORZOI SEQ EXTRACTION (I/O) in background.
    3. SPLICING — Batched GPU prediction, per-mutation Oncosplice assembly.
    4. TIS — Cross-mutation mega-batch via ``compute_tis_shift_metrics_batch``.
    5. Collect FOLDING results (already complete from background).
    6. TRANSCRIPTION — GPU inference on pre-extracted sequences.
    7. ASSEMBLE — Build MutationResult per mutation.
    """
    if max_gpu_batch is None:
        max_gpu_batch = 8

    n = len(mut_ids)
    if n == 0:
        return []

    timing: dict[str, float] = {}
    errors: list[dict[str, str]] = [{} for _ in range(n)]

    # ================================================================
    # Phase 1: PREPARE  (pre-warm gene cache, then per-mutation prep)
    # ================================================================
    _t_prepare = time.perf_counter()

    from ..variants import MutationalEvent
    events = []
    for mid in mut_ids:
        m = MutationalEvent(mid)
        if not m.compatible():
            raise ValueError(f"Mutations in event are incompatible: {mid}")
        events.append(m)

    # Pre-warm the shared gene cache in parallel (deduped by gene name).
    # _prepare_mutation will then hit the cache on every call.
    gene_names = list({m.gene for m in events})
    with ThreadPoolExecutor(max_workers=min(n_io_workers, len(gene_names))) as pool:
        list(pool.map(
            lambda name: _load_gene_with_timeout(name, organism, gene_load_timeout),
            gene_names,
        ))

    # Per-mutation prep using shared helper
    prep: list[dict] = []
    for mid in mut_ids:
        p = _prepare_mutation(mid, organism, transcript_id, context_bp, gene_load_timeout)
        prep.append(p)

    timing["prepare"] = round(time.perf_counter() - _t_prepare, 4)

    # ================================================================
    # Phase 2: Launch FOLDING (CPU) and BORZOI SEQ EXTRACTION (I/O)
    #           in background, then run GPU phases
    # ================================================================

    folding_results: list[FoldingResult | None] = [None] * n
    folding_futures = None
    fold_pool = None

    if run_folding:
        fold_args = [(i, p) for i, p in enumerate(prep) if p["status"] == "ok"]
        if fold_args:
            def _fold_one(args):
                i, p = args
                try:
                    return i, _run_folding_single(
                        p["ref_transcript"], p["event"], p["mut_id"], folding_window_size
                    )
                except Exception as e:
                    return i, FoldingResult(status="error", reason=str(e), mut_id=p["mut_id"])

            fold_pool = ThreadPoolExecutor(max_workers=min(n_folding_workers, len(fold_args)))
            folding_futures = [fold_pool.submit(_fold_one, args) for args in fold_args]

    # Pre-extract Borzoi sequences in background (I/O bound, overlaps with GPU)
    transcription_results: list[TranscriptionResult | None] = [None] * n
    trans_eligible: list[int] = []
    trans_seq_pairs = None
    io_pool = None

    if run_transcription:
        from ..transcription.prediction import _extract_sequences as _borzoi_extract

        trans_specs_raw: list[dict] = []
        for i, p in enumerate(prep):
            if p["status"] != "ok":
                continue
            m = p["event"]
            trans_eligible.append(i)
            trans_specs_raw.append({
                "chrom": str(m.chrom),
                "pos": p["central_pos"],
                "mutations": m.mutation_args(),
                "organism": organism,
            })

        if trans_specs_raw:
            def _extract_one(spec):
                return _borzoi_extract(
                    spec["chrom"], spec["pos"], spec["mutations"],
                    spec.get("organism", "hg38"),
                )

            io_pool = ThreadPoolExecutor(max_workers=min(n_io_workers, len(trans_specs_raw)))
            trans_seq_future = io_pool.map(_extract_one, trans_specs_raw)

    # ================================================================
    # Phase 3: SPLICING (GPU)
    # ================================================================
    _t_splicing = time.perf_counter()

    splicing_results: list[SplicingResult] = [SplicingResult() for _ in range(n)]

    if run_splicing:
        from ..splicing.engines import (
            _prepare_splice_input,
            _assemble_splice_df,
            batch_run_splicing_engine,
            adjoin_splicing_outcomes as _adjoin,
        )

        splice_jobs: list[dict] = []
        tl_map: dict[int, TranscriptLibrary] = {}

        for i, p in enumerate(prep):
            if p["status"] != "ok":
                continue
            try:
                m = p["event"]
                # Clone the raw selected transcript (lighter than ref_transcript)
                ref_for_splicing = p["selected"].clone()
                ref_for_splicing.generate_pre_mrna(
                    upstream=SPLICING_CONTEXT_BP,
                    downstream=SPLICING_CONTEXT_BP,
                )
                tl = TranscriptLibrary(ref_for_splicing, m, individual_mutations=False)
                tl_map[i] = tl
                for key, t in tl:
                    sp = _prepare_splice_input(t.pre_mrna, p["central_pos"])
                    splice_jobs.append({"mut_idx": i, "tl_key": key, "prep": sp})
            except Exception as e:
                logger.debug("Batch splicing prep failed for mutation %d: %s", i, e)
                errors[i]["splicing"] = str(e)

        if splice_jobs:
            all_seqs = [j["prep"]["padded_seq"] for j in splice_jobs]
            try:
                all_preds = batch_run_splicing_engine(all_seqs, engine=splicing_engine)
            except Exception as e:
                logger.warning("Batch splicing inference failed, falling back to sequential: %s", e)
                all_preds = []
                for seq in all_seqs:
                    try:
                        from ..splicing.engines import run_splicing_engine as _rse
                        all_preds.append(_rse(seq=seq, engine=splicing_engine))
                    except Exception as e2:
                        logger.debug("Sequential splicing fallback failed: %s", e2)
                        all_preds.append(([], []))

            from collections import defaultdict
            import pandas as pd
            per_mut_splicing: dict[int, dict[str, pd.DataFrame]] = defaultdict(dict)

            for job, (donor_probs, acceptor_probs) in zip(splice_jobs, all_preds):
                try:
                    df = _assemble_splice_df(donor_probs, acceptor_probs, job["prep"])
                    per_mut_splicing[job["mut_idx"]][job["tl_key"]] = df
                except Exception as e:
                    logger.debug("Splice df assembly failed for job %s: %s", job["tl_key"], e)

            splice_index_cache: dict[str, dict] = {}
            for i, preds_dict in per_mut_splicing.items():
                if i not in tl_map:
                    continue
                tl = tl_map[i]
                p = prep[i]
                m = p["event"]
                ref_transcript = p["ref_transcript"]
                central_pos = p["central_pos"]

                try:
                    tl.splicing_results = _adjoin(preds_dict, tl.ref)
                    splicing_preds = tl.get_event_columns("event")

                    # Clear scoped pre-mRNA so isoform clones use direct-exon assembly
                    tl.event._pre_mrna = None

                    ss = SpliceSimulator(
                        splicing_preds,
                        tl.event,
                        feature="event",
                        max_distance=SPLICE_MAX_DISTANCE,
                        min_total_prob=min_total_prob,
                        min_delta_magnitude=min_delta_magnitude,
                        mutations=list(m.mutation_args()),
                    )

                    base_report = _splicing_base_report(p["mut_id"], m, ref_transcript, splicing_engine, central_pos)
                    ss_metadata = ss.report(central_pos)
                    gene = p["gene"]
                    if m.gene not in splice_index_cache:
                        splice_index_cache[m.gene] = _build_splice_site_index(gene) if gene else {}
                    splicing_results[i] = _oncosplice_dataframe(ref_transcript, ss, ss_metadata, base_report, max_isoforms=max_isoforms, splice_index=splice_index_cache[m.gene])
                except Exception as e:
                    logger.debug("Batch oncosplice assembly failed for mutation %d: %s", i, e)
                    errors[i]["splicing"] = str(e)

    timing["splicing"] = round(time.perf_counter() - _t_splicing, 4)

    # ================================================================
    # Phase 4: TIS (cross-mutation mega-batch, GPU)
    # ================================================================
    _t_tis = time.perf_counter()

    tis_results: list[TISResult | None] = [None] * n

    if run_tis:
        from ..tis.metrics import compute_tis_shift_metrics_batch

        tis_eligible: list[int] = []
        tis_pairs: list[tuple[str, str]] = []
        tis_meta: list[dict] = []

        for i, p in enumerate(prep):
            if p["status"] != "ok":
                continue

            tis_inputs = _prepare_tis(p["ref_transcript"], p["event"], p["mut_id"], tis_proximity_window)
            if isinstance(tis_inputs, TISResult):
                tis_results[i] = tis_inputs
                continue

            ref_mrna_seq, var_mrna_seq, mrna_positions = tis_inputs
            tis_eligible.append(i)
            tis_pairs.append((ref_mrna_seq, var_mrna_seq))
            tis_meta.append({
                "mut_id": p["mut_id"],
                "gene": p["event"].gene,
                "transcript_id": p["ref_transcript"].transcript_id,
                "ref_mrna_seq": ref_mrna_seq,
                "var_mrna_seq": var_mrna_seq,
                "mrna_positions": mrna_positions,
            })

        if tis_pairs:
            try:
                batch_results = compute_tis_shift_metrics_batch(
                    tis_pairs, probability_threshold=tis_probability_threshold,
                )
                for j, (idx, result) in enumerate(zip(tis_eligible, batch_results)):
                    meta = tis_meta[j]
                    tis_results[idx] = _tis_completed_dict(
                        result, meta["mut_id"], meta["gene"], meta["transcript_id"],
                        meta["ref_mrna_seq"], meta["var_mrna_seq"], meta["mrna_positions"],
                    )
            except Exception as e:
                for j, idx in enumerate(tis_eligible):
                    meta = tis_meta[j]
                    tis_results[idx] = TISResult(
                        status="error", score=0.0, reason=str(e),
                        mut_id=meta["mut_id"], gene=meta["gene"])
                    errors[idx]["tis"] = str(e)

    timing["tis"] = round(time.perf_counter() - _t_tis, 4)

    # ================================================================
    # Collect FOLDING results (should be done by now — ran during GPU phases)
    # ================================================================
    _t_folding = time.perf_counter()

    if folding_futures is not None:
        for fut in folding_futures:
            idx, result = fut.result()
            folding_results[idx] = result
    if fold_pool is not None:
        fold_pool.shutdown(wait=False)

    timing["folding"] = round(time.perf_counter() - _t_folding, 4)

    # ================================================================
    # Phase 5: TRANSCRIPTION (GPU — sequences pre-extracted during GPU phases)
    # ================================================================
    _t_trans = time.perf_counter()

    if run_transcription and trans_eligible:
        try:
            trans_seq_pairs = list(trans_seq_future)
        except Exception as e:
            logger.warning("Borzoi sequence extraction failed: %s", e)
            trans_seq_pairs = None
    if io_pool is not None:
        io_pool.shutdown(wait=False)

    if run_transcription and trans_eligible and trans_seq_pairs:
        from ..transcription.model import _load_borzoi
        from ..transcription.prediction import (
            _predict_tracks_batch,
            _compute_tissue_changes,
            _build_gtex_mapping,
            _normalize_score,
        )

        try:
            model, device, dtype, use_autocast = _load_borzoi()
            gtex_map = _build_gtex_mapping()

            # Flatten all WT+mut sequences for batched inference
            all_seqs: list[str] = []
            for wt, mut in trans_seq_pairs:
                all_seqs.append(wt)
                all_seqs.append(mut)

            # Chunked GPU forward passes
            all_tracks: list[np.ndarray] = []
            for i in range(0, len(all_seqs), max_gpu_batch):
                chunk = all_seqs[i : i + max_gpu_batch]
                all_tracks.extend(
                    _predict_tracks_batch(model, chunk, device, dtype, use_autocast)
                )

            # Compute per-mutation tissue changes
            for j, idx in enumerate(trans_eligible):
                p = prep[idx]
                wt_tracks = all_tracks[j * 2]
                mut_tracks = all_tracks[j * 2 + 1]
                tissue_data = _compute_tissue_changes(wt_tracks, mut_tracks, gtex_map)
                score = _normalize_score(tissue_data["max_abs_lfc"])
                tissue_data["regulatory_score"] = round(score, 4)
                tissue_data["status"] = "completed"
                transcription_results[idx] = _transcription_result_from_dict(
                    tissue_data, p["mut_id"], p["event"].gene)
        except Exception as e:
            for idx in trans_eligible:
                p = prep[idx]
                transcription_results[idx] = TranscriptionResult(
                    status="error", reason=str(e),
                    mut_id=p["mut_id"], gene=p["event"].gene)
                errors[idx]["borzoi"] = str(e)

    timing["borzoi"] = round(time.perf_counter() - _t_trans, 4)
    timing["total"] = round(time.perf_counter() - _t_prepare, 4)

    # ================================================================
    # Phase 6: ASSEMBLE
    # ================================================================

    results: list[MutationResult] = []
    for i in range(n):
        p = prep[i]
        if p["status"] != "ok":
            results.append(MutationResult(timing=timing, errors=errors[i]))
        else:
            results.append(MutationResult(
                splicing=splicing_results[i],
                tis=tis_results[i],
                folding=folding_results[i],
                transcription=transcription_results[i],
                transcript=p["mut_transcript"],
                mutation_context=p["mutation_context"],
                timing=timing,
                errors=errors[i],
            ))

    return results
