# geney/concurrent.py — Concurrent mutation analysis with shared GPU batch queues.
"""
Multi-worker mutation analysis that maximises GPU utilisation.

Architecture::

    Worker threads (N)              GPU batch queues (per model)
    ┌──────────────┐               ┌───────────────────────┐
    │ gene load    │── splice req ─▶│ accumulate → batch    │
    │ assembly     │◀── splice res ─│ forward → scatter     │
    │ oncosplice   │── tis req ────▶│                       │
    │ folding      │◀── tis res ───│                       │
    │ context prep │── borzoi req ─▶│                       │
    │              │◀── borzoi res ─│                       │
    └──────────────┘               └───────────────────────┘

Each worker runs the full CPU pipeline for one mutation. When it
needs GPU inference it submits work to a shared queue and blocks
on a Future. The queue thread accumulates items from all workers
and dispatches batched forward passes, splitting results back.

Usage::

    from geney import Geney

    g = Geney(device="cuda")
    results = g.analyze_concurrent(
        ["KRAS:12:25227343:G:T", "TP53:17:7674220:C:T", ...],
        n_workers=8,
        max_gpu_batch=16,
    )
"""
from __future__ import annotations

import logging
import threading
import time
from collections import deque
from concurrent.futures import Future, ThreadPoolExecutor
from typing import Any

import numpy as np

from .pipelines import SPLICING_CONTEXT_BP, SPLICE_MAX_DISTANCE

from .results import (
    SplicingResult,
    ProteinResult,
    TISResult,
    FoldingResult,
    TranscriptionResult,
    MiRNAResult,
    MutationResult,
)

logger = logging.getLogger(__name__)


# ============================================================================
# GPU Batch Queue
# ============================================================================

class GPUBatchQueue:
    """Thread-safe queue that accumulates GPU work and dispatches in batches.

    Workers call :meth:`submit` with a list of items (one mutation's worth
    of GPU inputs).  The internal dispatcher thread collects items from all
    workers and fires a single batched ``dispatch_fn`` call when either
    *max_batch* items accumulate or *timeout_s* elapses — whichever comes
    first.

    Parameters
    ----------
    dispatch_fn : callable
        ``fn(list[T]) -> list[R]``.  Called with flattened inputs from
        multiple workers; must return one result per input in order.
    max_batch : int
        Fire as soon as this many items accumulate.
    timeout_s : float
        Maximum seconds to wait before firing an under-full batch.
    """

    def __init__(
        self,
        dispatch_fn,
        max_batch: int = 16,
        timeout_s: float = 0.05,
    ):
        self._dispatch_fn = dispatch_fn
        self._max_batch = max_batch
        self._timeout = timeout_s

        self._pending: deque[tuple[list, Future]] = deque()
        self._total_items = 0
        self._lock = threading.Lock()
        self._ready = threading.Event()
        self._shutdown = False

        self._thread = threading.Thread(target=self._dispatcher, daemon=True)
        self._thread.start()

    def submit(self, items: list) -> Future:
        """Submit a group of items.  Returns Future resolving to list[result]."""
        fut: Future = Future()
        if not items:
            fut.set_result([])
            return fut

        with self._lock:
            self._pending.append((items, fut))
            self._total_items += len(items)
            if self._total_items >= self._max_batch:
                self._ready.set()

        # Also signal so the timeout path wakes up
        self._ready.set()
        return fut

    def _dispatcher(self):
        """Background thread: waits, collects, dispatches, scatters."""
        while not self._shutdown:
            # Wait until signalled or timeout
            self._ready.wait(timeout=self._timeout)
            self._ready.clear()

            # Grab all pending work
            with self._lock:
                if not self._pending:
                    continue
                batch = list(self._pending)
                self._pending.clear()
                self._total_items = 0

            # Flatten items from all workers
            flat_items: list = []
            groups: list[tuple[Future, int, int]] = []
            for items, fut in batch:
                groups.append((fut, len(flat_items), len(items)))
                flat_items.extend(items)

            # Dispatch (GPU forward pass)
            try:
                flat_results = self._dispatch_fn(flat_items)
                for fut, start, count in groups:
                    fut.set_result(flat_results[start : start + count])
            except Exception as e:
                for fut, _, _ in groups:
                    if not fut.done():
                        fut.set_exception(e)

    def shutdown(self):
        """Signal dispatcher to stop and wait for it to finish."""
        self._shutdown = True
        self._ready.set()
        self._thread.join(timeout=5)


# ============================================================================
# Dispatch functions (thin wrappers around existing batch inference)
# ============================================================================

def _make_splice_dispatcher(engine: str):
    """Return a dispatch function for batched splicing inference."""
    from .splicing.engines import batch_run_splicing_engine

    def dispatch(sequences: list[str]) -> list[tuple]:
        return batch_run_splicing_engine(sequences, engine=engine)

    return dispatch


def _make_tis_dispatcher():
    """Return a dispatch function for batched TIS context scoring."""
    from .tis.prediction import predict_context_scores_batch

    def dispatch(contexts: list[str]) -> list[float]:
        scores = predict_context_scores_batch(contexts)
        return scores.tolist()

    return dispatch


def _make_borzoi_dispatcher():
    """Return a dispatch function for batched Borzoi track prediction."""
    from .transcription.model import _load_borzoi
    from .transcription.prediction import _predict_tracks_batch

    model, device, dtype, use_autocast = _load_borzoi()

    def dispatch(sequences: list[str]) -> list[np.ndarray]:
        return _predict_tracks_batch(model, sequences, device, dtype, use_autocast)

    return dispatch


# ============================================================================
# Worker function (one mutation)
# ============================================================================

def _concurrent_worker(
    mut_id: str,
    *,
    organism: str,
    transcript_id: str | None,
    splicing_engine: str,
    context_bp: int,
    run_splicing: bool,
    run_all_transcripts: bool,
    run_tis: bool,
    run_folding: bool,
    run_transcription: bool,
    run_mirna: bool,
    mirna_seeds: dict[str, str] | None,
    mirna_full_prediction: bool,
    tis_proximity_window: int,
    tis_probability_threshold: float,
    folding_window_size: int,
    folding_compute_null: bool,
    gene_load_timeout: float | None,
    max_isoforms: int | None,
    min_total_prob: float,
    min_delta_magnitude: float,
    splice_queue: GPUBatchQueue | None,
    borzoi_queue: GPUBatchQueue | None,
    tis_lock: threading.Lock,
) -> MutationResult:
    """Run the full mutation pipeline for one mutation, using GPU queues."""
    from .pipelines import (
        _prepare_mutation,
        _splicing_base_report,
        _oncosplice_dataframe,
        _build_splice_site_index,
        _check_branch_point,
        _get_conservation_vector,
        _prepare_tis,
        _tis_completed_dict,
        _run_folding_single,
        _transcription_result_from_dict,
    )
    from .splicing import SpliceSimulator, TranscriptLibrary
    from .splicing.engines import (
        _prepare_splice_input,
        _assemble_splice_df,
        adjoin_splicing_outcomes as _adjoin,
    )

    timing: dict[str, float] = {}
    errors: dict[str, str] = {}
    _t_total = time.perf_counter()

    # -- Prepare (parse → gene → transcript → mutate → context) --
    _t0 = time.perf_counter()
    p = _prepare_mutation(mut_id, organism, transcript_id, context_bp, gene_load_timeout)
    timing["prepare"] = round(time.perf_counter() - _t0, 4)

    if p["status"] != "ok":
        timing["total"] = round(time.perf_counter() - _t_total, 4)
        errors["prepare"] = p["status"]
        return MutationResult(
            splicing=SplicingResult(status=p["status"],
                                   reason=f"Prepare failed: {p['status']}"),
            timing=timing,
            errors=errors,
        )

    m = p["event"]
    central_pos = p["central_pos"]
    gene = p["gene"]
    selected = p["selected"]
    reference_transcript = p["ref_transcript"]
    mutated_transcript = p["mut_transcript"]
    mutation_context = p["mutation_context"]

    # -- 1. Splicing (GPU via queue) --
    splicing_result = SplicingResult()
    _shared_splicing_preds = None
    if run_splicing and splice_queue is not None:
        _t0 = time.perf_counter()
        try:
            ref_for_splicing = 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)

            # Prepare splice inputs (CPU)
            splice_preps = {}
            seqs_to_submit = []
            keys_order = []
            for key, t in tl:
                sp = _prepare_splice_input(t.pre_mrna, central_pos)
                splice_preps[key] = sp
                seqs_to_submit.append(sp["padded_seq"])
                keys_order.append(key)

            # Submit to GPU queue and block
            future = splice_queue.submit(seqs_to_submit)
            predictions = future.result()

            # Assemble splice DataFrames (CPU)
            import pandas as pd
            splice_dfs = {}
            for key, sp, (don_probs, acc_probs) in zip(keys_order, [splice_preps[k] for k in keys_order], predictions):
                splice_dfs[key] = _assemble_splice_df(don_probs, acc_probs, sp)

            tl.splicing_results = _adjoin(splice_dfs, tl.ref)
            splicing_preds = tl.get_event_columns("event")
            _shared_splicing_preds = splicing_preds

            # Clear scoped pre-mRNA for 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(mut_id, m, reference_transcript, splicing_engine, central_pos)
            ss_metadata = ss.report(central_pos)
            splice_index = _build_splice_site_index(gene)
            splicing_result = _oncosplice_dataframe(
                reference_transcript, ss, ss_metadata, base_report,
                max_isoforms=max_isoforms, splice_index=splice_index,
            )
            # Branch point check
            bp_disrupted, bp_info = _check_branch_point(
                reference_transcript, m.positions, ref_for_splicing.pre_mrna)
            splicing_result.branch_point_disrupted = bp_disrupted
            splicing_result.branch_point_info = bp_info
        except Exception as e:
            splicing_result = SplicingResult()
            errors["splicing"] = str(e)
        timing["spliceai"] = round(time.perf_counter() - _t0, 4)

    # -- 1a. Multi-transcript splicing (reuses shared SpliceAI predictions) --
    multi_transcript_result = None
    if run_all_transcripts and _shared_splicing_preds is not None:
        _t0 = time.perf_counter()
        try:
            from .pipelines._multi_transcript import analyze_all_transcripts
            multi_transcript_result = analyze_all_transcripts(
                gene=gene,
                mutation_event=m,
                splicing_preds=_shared_splicing_preds,
                primary_transcript=reference_transcript,
                primary_splicing_result=splicing_result,
                splicing_engine=splicing_engine,
                central_pos=central_pos,
                max_isoforms=max_isoforms,
                min_total_prob=min_total_prob,
                min_delta_magnitude=min_delta_magnitude,
            )
        except Exception as e:
            logger.debug("Multi-transcript analysis failed for %s: %s", mut_id, e)
            errors["multi_transcript"] = str(e)
        timing["multi_transcript"] = round(time.perf_counter() - _t0, 4)

    # -- 1b. Protein-level analysis (CPU, fast) --
    protein_result = None
    _t0 = time.perf_counter()
    try:
        ref_prot = reference_transcript.protein or ""
        var_prot = getattr(mutated_transcript, "protein", "") or ""
        if ref_prot:
            from .protein import analyze_protein_change
            cons = _get_conservation_vector(reference_transcript)
            protein_result = ProteinResult(
                **{k: v for k, v in analyze_protein_change(
                    ref_prot, var_prot, cons,
                    mut_id=mut_id, gene=m.gene,
                    transcript_id=reference_transcript.transcript_id,
                ).__dict__.items() if k in ProteinResult.__dataclass_fields__}
            )
    except Exception as e:
        protein_result = ProteinResult(status="error", reason=str(e),
                                       mut_id=mut_id, gene=m.gene)
        errors["protein"] = str(e)
    timing["protein"] = round(time.perf_counter() - _t0, 4)

    # -- 2. TIS (serialised with lock — fast enough) --
    tis_result = None
    if run_tis:
        _t0 = time.perf_counter()
        try:
            from .tis import compute_tis_shift_metrics

            tis_inputs = _prepare_tis(reference_transcript, m, mut_id, tis_proximity_window)
            if isinstance(tis_inputs, TISResult):
                tis_result = tis_inputs
            else:
                ref_mrna_seq, var_mrna_seq, mrna_positions = tis_inputs
                with tis_lock:
                    result = compute_tis_shift_metrics(
                        ref_mrna=ref_mrna_seq,
                        var_mrna=var_mrna_seq,
                        canonical_start=None,
                        probability_threshold=tis_probability_threshold,
                    )
                tis_result = _tis_completed_dict(
                    result, mut_id, m.gene, reference_transcript.transcript_id,
                    ref_mrna_seq, var_mrna_seq, mrna_positions,
                )
        except Exception as e:
            tis_result = TISResult(status="error", score=0.0, reason=str(e),
                                   mut_id=mut_id, gene=m.gene)
            errors["tis"] = str(e)
        timing["tis"] = round(time.perf_counter() - _t0, 4)

    # -- 3. Folding (CPU only) --
    folding_result = None
    if run_folding:
        _t0 = time.perf_counter()
        try:
            folding_result = _run_folding_single(reference_transcript, m, mut_id,
                                                 folding_window_size,
                                                 compute_null=folding_compute_null)
        except Exception as e:
            folding_result = FoldingResult(status="error", reason=str(e),
                                           mut_id=mut_id, gene=m.gene)
            errors["folding"] = str(e)
        timing["folding"] = round(time.perf_counter() - _t0, 4)

    # -- 4. Transcription / Borzoi (GPU via queue) --
    transcription_result = None
    if run_transcription and borzoi_queue is not None:
        _t0 = time.perf_counter()
        try:
            from .transcription.prediction import (
                _extract_sequences,
                _compute_tissue_changes,
                _build_gtex_mapping,
                _normalize_score,
            )

            # Extract WT + mutant sequences (I/O)
            wt_seq, mut_seq = _extract_sequences(
                chrom=str(m.chrom), pos=central_pos,
                mutations=m.mutation_args(), organism=organism,
            )

            # Submit to GPU queue and block
            future = borzoi_queue.submit([wt_seq, mut_seq])
            tracks = future.result()  # [wt_tracks, mut_tracks]

            # Post-process (CPU)
            gtex_map = _build_gtex_mapping()
            tissue_data = _compute_tissue_changes(tracks[0], tracks[1], gtex_map)
            score = _normalize_score(tissue_data["max_abs_lfc"])
            tissue_data["regulatory_score"] = round(score, 4)
            tissue_data["status"] = "completed"
            transcription_result = _transcription_result_from_dict(tissue_data, mut_id, m.gene)
        except Exception as e:
            transcription_result = TranscriptionResult(
                status="error", reason=str(e), mut_id=mut_id, gene=m.gene)
            errors["borzoi"] = str(e)
        timing["borzoi"] = round(time.perf_counter() - _t0, 4)

    # -- 5. miRNA binding site disruption (CPU only) --
    mirna_result = None
    if run_mirna and mirna_seeds:
        _t0 = time.perf_counter()
        try:
            from .mirna import analyze_mirna_disruption

            ref_mrna_seq = getattr(reference_transcript.mature_mrna, "seq", "")
            var_mrna_obj = getattr(mutated_transcript, "mature_mrna", None)
            var_mrna_seq = getattr(var_mrna_obj, "seq", "") if var_mrna_obj else ""

            if ref_mrna_seq and var_mrna_seq:
                cds_len = len(reference_transcript.protein or "") * 3
                utr3_start = cds_len if cds_len > 0 else None

                result = analyze_mirna_disruption(
                    ref_mrna=ref_mrna_seq,
                    var_mrna=var_mrna_seq,
                    mirna_seeds=mirna_seeds,
                    utr3_start=utr3_start,
                    mut_id=mut_id,
                    gene=m.gene,
                )

                mibsim_repression_ref = None
                mibsim_repression_var = None
                has_changes = result.n_gained + result.n_lost + result.n_modified > 0
                if mirna_full_prediction and has_changes:
                    try:
                        from .mirna.features import compute_site_features
                        from .mirna.mibsim import predict_transcript_repression

                        utr5 = getattr(reference_transcript, "utr5_seq",
                                getattr(reference_transcript, "utr5", "")) or ""
                        orf_seq = getattr(reference_transcript, "cds_seq",
                                  getattr(reference_transcript, "cds", "")) or ""
                        utr3_seq = getattr(reference_transcript, "utr3_seq",
                                   getattr(reference_transcript, "utr3", "")) or ""

                        if not orf_seq and ref_mrna_seq:
                            prot_len = len(reference_transcript.protein or "")
                            orf_len = prot_len * 3
                            utr5 = utr5 or ref_mrna_seq[:max(0, len(ref_mrna_seq) - orf_len)]
                            orf_seq = ref_mrna_seq[len(utr5):len(utr5) + orf_len]
                            utr3_seq = utr3_seq or ref_mrna_seq[len(utr5) + orf_len:]

                        if orf_seq:
                            cons_vec = _get_conservation_vector(reference_transcript)
                            has_real_cons = cons_vec is not None and not (
                                hasattr(cons_vec, 'size') and cons_vec.size > 0
                                and np.allclose(cons_vec, 1.0))

                            def _build_site_inputs(sites):
                                inputs = []
                                for site in sites:
                                    mir_seq = mirna_seeds.get(site.mirna_id, "")
                                    if not mir_seq:
                                        continue
                                    feats = compute_site_features(
                                        utr5=utr5, orf=orf_seq, utr3=utr3_seq,
                                        mirna=mir_seq, rna_start=site.position,
                                        seed_length=site.end - site.position,
                                        seed_type=site.seed_type,
                                        region="UTR3" if site.region == "utr3" else "ORF",
                                        with_conservation=has_real_cons, with_thermo=False)
                                    inputs.append({"features": feats, "position": site.position,
                                                   "region": site.region, "seed_type": site.seed_type})
                                return inputs

                            if result.ref_sites:
                                ref_inputs = _build_site_inputs(result.ref_sites)
                                if ref_inputs:
                                    mibsim_repression_ref = predict_transcript_repression(ref_inputs)["total_repression"]
                            if result.var_sites:
                                var_inputs = _build_site_inputs(result.var_sites)
                                if var_inputs:
                                    mibsim_repression_var = predict_transcript_repression(var_inputs)["total_repression"]

                    except Exception as e:
                        logger.debug("miBSIM full prediction failed for %s: %s", mut_id, e)

                mirna_result = MiRNAResult(
                    status=result.status,
                    score=result.score,
                    n_sites_ref=result.n_sites_ref,
                    n_sites_var=result.n_sites_var,
                    n_gained=result.n_gained,
                    n_lost=result.n_lost,
                    n_modified=result.n_modified,
                    n_unchanged=result.n_unchanged,
                    total_repression_ref=mibsim_repression_ref if mibsim_repression_ref is not None else result.total_repression_ref,
                    total_repression_var=mibsim_repression_var if mibsim_repression_var is not None else result.total_repression_var,
                    repression_delta=round(
                        (mibsim_repression_var or result.total_repression_var)
                        - (mibsim_repression_ref or result.total_repression_ref), 4),
                    top_affected_mirnas=result.top_affected_mirnas,
                    site_changes=result.site_changes,
                    n_mirnas_scanned=result.n_mirnas_scanned,
                    mut_id=mut_id,
                    gene=m.gene,
                )
            else:
                mirna_result = MiRNAResult(
                    status="skipped", reason="no_mature_mrna",
                    mut_id=mut_id, gene=m.gene)
        except Exception as e:
            mirna_result = MiRNAResult(
                status="error", reason=str(e), mut_id=mut_id, gene=m.gene)
            errors["mirna"] = str(e)
        timing["mirna"] = round(time.perf_counter() - _t0, 4)

    timing["total"] = round(time.perf_counter() - _t_total, 4)

    return MutationResult(
        splicing=splicing_result,
        multi_transcript=multi_transcript_result,
        protein=protein_result,
        tis=tis_result,
        folding=folding_result,
        transcription=transcription_result,
        mirna=mirna_result,
        transcript=mutated_transcript,
        mutation_context=mutation_context,
        timing=timing,
        errors=errors,
    )


# ============================================================================
# Public entry point
# ============================================================================

def analyze_concurrent(
    mut_ids: list[str],
    *,
    organism: str = "hg38",
    splicing_engine: str = "ensemble",
    transcript_id: str | None = None,
    context_bp: int = 1000,
    run_splicing: bool = True,
    run_all_transcripts: bool = False,
    run_tis: bool = True,
    run_folding: bool = True,
    run_transcription: bool = True,
    run_mirna: bool = False,
    mirna_seeds: dict[str, str] | None = None,
    mirna_full_prediction: bool = True,
    tis_proximity_window: int = 500,
    tis_probability_threshold: float = 0.1,
    folding_window_size: int = 39,
    folding_compute_null: bool = False,
    gene_load_timeout: float | None = None,
    max_isoforms: int | None = None,
    min_total_prob: float = 0.5,
    min_delta_magnitude: float = 0.2,
    n_workers: int = 8,
    max_gpu_batch: int = 16,
    batch_timeout_ms: int = 50,
) -> list[MutationResult]:
    """Analyse mutations concurrently with shared GPU batch queues.

    Spawns *n_workers* threads that each run the full mutation pipeline.
    CPU-bound work (gene loading, transcript assembly, oncosplice scoring,
    ViennaRNA folding) runs in parallel across workers.  GPU inference
    (splicing, TIS, Borzoi) is funnelled through batch queues so the GPU
    sees large batches even though workers submit one-at-a-time.

    Parameters
    ----------
    mut_ids : list[str]
        Mutation IDs to analyse.
    n_workers : int
        Number of concurrent worker threads (default 8).
    max_gpu_batch : int
        Maximum items per GPU dispatch (default 16).
    batch_timeout_ms : int
        Milliseconds to wait before dispatching an under-full batch
        (default 50).
    **kwargs
        All other parameters are forwarded to the per-mutation pipeline
        (same as :meth:`Geney.analyze`).

    Returns
    -------
    list[MutationResult]
        One result per input mutation, in the same order.
    """
    if not mut_ids:
        return []

    timeout_s = batch_timeout_ms / 1000.0
    queues: list[GPUBatchQueue] = []

    # -- Set up GPU queues --
    splice_queue = None
    if run_splicing:
        splice_queue = GPUBatchQueue(
            _make_splice_dispatcher(splicing_engine),
            max_batch=max_gpu_batch,
            timeout_s=timeout_s,
        )
        queues.append(splice_queue)

    borzoi_queue = None
    if run_transcription:
        try:
            borzoi_queue = GPUBatchQueue(
                _make_borzoi_dispatcher(),
                max_batch=max_gpu_batch,
                timeout_s=timeout_s,
            )
            queues.append(borzoi_queue)
        except Exception as e:
            logger.warning("Borzoi queue init failed (will skip transcription): %s", e)
            run_transcription = False

    tis_lock = threading.Lock()

    # -- Load miRNA seeds once (shared across all workers) --
    if run_mirna and not mirna_seeds:
        from .mirna.reference import get_mirna_panel
        mirna_seeds = get_mirna_panel()

    # -- Shared kwargs for workers --
    worker_kwargs = dict(
        organism=organism,
        transcript_id=transcript_id,
        splicing_engine=splicing_engine,
        context_bp=context_bp,
        run_splicing=run_splicing,
        run_all_transcripts=run_all_transcripts,
        run_tis=run_tis,
        run_folding=run_folding,
        run_transcription=run_transcription,
        run_mirna=run_mirna,
        mirna_seeds=mirna_seeds,
        mirna_full_prediction=mirna_full_prediction,
        tis_proximity_window=tis_proximity_window,
        tis_probability_threshold=tis_probability_threshold,
        folding_window_size=folding_window_size,
        folding_compute_null=folding_compute_null,
        gene_load_timeout=gene_load_timeout,
        max_isoforms=max_isoforms,
        min_total_prob=min_total_prob,
        min_delta_magnitude=min_delta_magnitude,
        splice_queue=splice_queue,
        borzoi_queue=borzoi_queue,
        tis_lock=tis_lock,
    )

    # -- Run workers --
    t0 = time.perf_counter()
    results: list[MutationResult] = [MutationResult() for _ in range(len(mut_ids))]

    def _safe_worker(idx: int, mid: str) -> tuple[int, MutationResult]:
        try:
            return idx, _concurrent_worker(mid, **worker_kwargs)
        except Exception as e:
            return idx, MutationResult(errors={"worker": str(e)})

    with ThreadPoolExecutor(max_workers=n_workers) as pool:
        futures = [
            pool.submit(_safe_worker, i, mid)
            for i, mid in enumerate(mut_ids)
        ]
        for fut in futures:
            idx, result = fut.result()
            results[idx] = result

    # -- Shutdown queues --
    for q in queues:
        q.shutdown()

    elapsed = time.perf_counter() - t0
    logger.info(
        "analyze_concurrent: %d mutations in %.2fs (%.2f mut/s, %d workers)",
        len(mut_ids), elapsed, len(mut_ids) / elapsed if elapsed > 0 else 0, n_workers,
    )

    return results
