# 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 .results import (
    SplicingResult,
    TISResult,
    FoldingResult,
    TranscriptionResult,
    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_tis: bool,
    run_folding: bool,
    run_transcription: bool,
    tis_proximity_window: int,
    tis_probability_threshold: float,
    folding_window_size: int,
    gene_load_timeout: float | None,
    max_isoforms: int | None,
    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 .variants import MutationalEvent
    from .utils import select_transcript, _apply_mutation_safe
    from .pipelines import (
        _load_gene_with_timeout,
        _splicing_base_report,
        _oncosplice_dataframe,
        _build_splice_site_index,
        _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()

    # -- Parse --
    _t0 = time.perf_counter()
    m = MutationalEvent(mut_id)
    if not m.compatible():
        raise ValueError(f"Mutations in event are incompatible: {mut_id}")
    central_pos = m.central_position
    timing["parse"] = round(time.perf_counter() - _t0, 4)

    # -- Gene load --
    _t0 = time.perf_counter()
    gene = _load_gene_with_timeout(m.gene, organism, gene_load_timeout)
    timing["gene_load"] = round(time.perf_counter() - _t0, 4)
    if gene is None:
        timing["total"] = round(time.perf_counter() - _t_total, 4)
        return MutationResult(timing=timing)

    # -- Transcript prep --
    _t0 = time.perf_counter()
    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        timing["transcript_prep"] = round(time.perf_counter() - _t0, 4)
        timing["total"] = round(time.perf_counter() - _t_total, 4)
        return MutationResult(timing=timing)

    reference_transcript = selected.generate_mature_mrna().generate_protein()
    timing["transcript_prep"] = round(time.perf_counter() - _t0, 4)

    # -- Mutate transcript --
    _t0 = time.perf_counter()
    mutated_transcript = selected.clone().generate_mature_mrna()
    for pos, ref, alt in m.mutation_args():
        try:
            _apply_mutation_safe(mutated_transcript.mature_mrna, pos, ref, alt)
        except Exception:
            pass

    # Scoped pre-mRNA for mutation context
    ctx_transcript = selected.clone()
    ctx_transcript.generate_pre_mrna(
        region_start=central_pos - context_bp,
        region_end=central_pos + context_bp,
    )
    for pos, ref, alt in m.mutation_args():
        try:
            _apply_mutation_safe(ctx_transcript.pre_mrna, pos, ref, alt)
        except Exception:
            pass
    mutation_context = ctx_transcript.pre_mrna

    mutated_transcript.generate_protein()
    timing["mutate_transcript"] = round(time.perf_counter() - _t0, 4)

    # -- 1. Splicing (GPU via queue) --
    splicing_result = SplicingResult()
    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(
                region_start=central_pos - 7500,
                region_end=central_pos + 7500,
            )
            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")

            # Clear scoped pre-mRNA for direct-exon assembly
            tl.event._pre_mrna = None

            ss = SpliceSimulator(
                splicing_preds, tl.event, feature="event", max_distance=100_000_000,
                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,
            )
        except Exception as e:
            splicing_result = SplicingResult()
            errors["splicing"] = str(e)
        timing["spliceai"] = round(time.perf_counter() - _t0, 4)

    timing["oncosplice"] = 0.0

    # -- 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)
        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)

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

    return MutationResult(
        splicing=splicing_result,
        tis=tis_result,
        folding=folding_result,
        transcription=transcription_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 = "spliceai",
    transcript_id: str | None = None,
    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_isoforms: int | None = None,
    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()

    # -- 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_tis=run_tis,
        run_folding=run_folding,
        run_transcription=run_transcription,
        tis_proximity_window=tis_proximity_window,
        tis_probability_threshold=tis_probability_threshold,
        folding_window_size=folding_window_size,
        gene_load_timeout=gene_load_timeout,
        max_isoforms=max_isoforms,
        splice_queue=splice_queue,
        borzoi_queue=borzoi_queue,
        tis_lock=tis_lock,
    )

    # -- Run workers --
    t0 = time.perf_counter()
    results: list[MutationResult] = [MutationResult()] * 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
