# geney/engine.py — Unified configuration and entry point.
"""
Geney engine: consolidates model paths, device, and pipeline parameters
into a single object so that init happens once and every subsequent
``analyze()`` call reuses the same configuration.

Usage::

    from geney import Geney

    g = Geney(
        titer_model_dir="/data/models/titer",
        openspliceai_model_dir="/data/models/openspliceai-mane/10000nt",
        device="cuda",                           # optional, auto-detects
        splicing_engine="ensemble",              # default engine
        organism="hg38",                         # default organism
    )

    result = g.analyze("KRAS:12:25227343:G:T")
"""
from __future__ import annotations

from .pipelines import mutation_pipeline, mutation_pipeline_batch, MutationResult


class Geney:
    """Unified geney configuration and analysis entry point.

    All model-path and device setters are called at construction time so that
    the first ``analyze()`` call doesn't pay a surprise configuration cost.

    Parameters
    ----------
    titer_model_dir : str, optional
        Directory containing TITER ``.pt`` / ``.hdf5`` model weights.
        Falls back to ``TITER_MODEL_DIR`` env var → bundled models.
    openspliceai_model_dir : str, optional
        Directory containing OpenSpliceAI MANE 10 000 nt model weights.
        Falls back to ``OPENSPLICEAI_MODEL_DIR`` env var → bundled models.
    borzoi_model : str, optional
        HuggingFace model name or local path for Borzoi transcription
        predictions.  Falls back to ``BORZOI_MODEL`` env var → default
        ``"johahi/borzoi-replicate-0"``.
    device : str, optional
        PyTorch device for TITER, Pangolin, and Borzoi (``"cpu"``, ``"cuda"``,
        ``"mps"``).  *None* → auto-detect per subsystem.
    splicing_engine : str
        Default splicing engine (``"spliceai"`` or ``"pangolin"``).
    organism : str
        Default genome build passed to seqmat (``"hg38"``).
    """

    def __init__(
        self,
        titer_model_dir: str | None = None,
        openspliceai_model_dir: str | None = None,
        borzoi_model: str | None = None,
        device: str | None = None,
        splicing_engine: str = "ensemble",
        organism: str = "hg38",
    ):
        self.splicing_engine = splicing_engine
        self.organism = organism

        # -- TITER models --
        if titer_model_dir is not None:
            from .tis import set_titer_model_dir
            set_titer_model_dir(titer_model_dir)

        # -- OpenSpliceAI models --
        if openspliceai_model_dir is not None:
            from .splicing import set_openspliceai_model_dir
            set_openspliceai_model_dir(openspliceai_model_dir)

        # -- Borzoi model --
        if borzoi_model is not None:
            from .transcription import set_borzoi_model
            set_borzoi_model(borzoi_model)

        # -- Device overrides --
        if device is not None:
            from .tis.model import set_titer_device
            from .splicing.engines import set_splicing_device
            from .transcription.model import set_borzoi_device
            set_titer_device(device)
            set_splicing_device(device)
            set_borzoi_device(device)

    def analyze(
        self,
        mut_id: str,
        transcript_id: str | None = None,
        *,
        organism: str | None = None,
        splicing_engine: 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,
        enable_borzoi: bool | None = None,
        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,
    ) -> MutationResult:
        """Run the unified mutation pipeline.

        Accepts the same parameters as :func:`mutation_pipeline` but defaults
        *organism* and *splicing_engine* to the values set at construction.

        Parameters
        ----------
        enable_borzoi : bool, optional
            Convenience alias for *run_transcription*.  When provided,
            overrides *run_transcription*.
        max_isoforms : int, optional
            Maximum number of splice isoforms to score (ranked by prevalence).
            *None* means no limit.
        """
        if enable_borzoi is not None:
            run_transcription = enable_borzoi
        return mutation_pipeline(
            mut_id=mut_id,
            transcript_id=transcript_id,
            organism=organism or self.organism,
            splicing_engine=splicing_engine or self.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,
        )

    def analyze_batch(
        self,
        mut_ids: list[str],
        transcript_id: str | None = None,
        *,
        organism: str | None = None,
        splicing_engine: 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_gpu_batch: int | None = None,
        n_folding_workers: int = 4,
        n_io_workers: int = 8,
        max_isoforms: int | None = None,
    ) -> list[MutationResult]:
        """Batch-optimized mutation analysis.

        Same as calling :meth:`analyze` N times but batches GPU inference
        across mutations.  See :func:`mutation_pipeline_batch` for details.
        """
        return mutation_pipeline_batch(
            mut_ids=mut_ids,
            transcript_id=transcript_id,
            organism=organism or self.organism,
            splicing_engine=splicing_engine or self.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_gpu_batch=max_gpu_batch,
            n_folding_workers=n_folding_workers,
            n_io_workers=n_io_workers,
            max_isoforms=max_isoforms,
        )

    def analyze_concurrent(
        self,
        mut_ids: list[str],
        transcript_id: str | None = None,
        *,
        organism: str | None = None,
        splicing_engine: 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,
        enable_borzoi: bool | None = None,
        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]:
        """Concurrent mutation analysis with shared GPU batch queues.

        Spawns *n_workers* threads that each run the full pipeline for one
        mutation.  CPU work (gene loading, transcript assembly, oncosplice,
        ViennaRNA) runs in parallel.  GPU inference (splicing, TIS, Borzoi)
        is funnelled through batch queues so the GPU sees large batches.

        Parameters
        ----------
        mut_ids : list[str]
            Mutation IDs to analyse.
        n_workers : int
            Concurrent CPU threads (default 8).
        max_gpu_batch : int
            Max items per GPU forward pass (default 16).
        batch_timeout_ms : int
            Max ms to wait before dispatching an under-full batch (default 50).
        enable_borzoi : bool, optional
            Convenience alias for *run_transcription*.
        **kwargs
            Same parameters as :meth:`analyze`.

        Returns
        -------
        list[MutationResult]
            One result per input mutation, in order.
        """
        from .concurrent import analyze_concurrent

        if enable_borzoi is not None:
            run_transcription = enable_borzoi
        return analyze_concurrent(
            mut_ids=mut_ids,
            organism=organism or self.organism,
            splicing_engine=splicing_engine or self.splicing_engine,
            transcript_id=transcript_id,
            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,
            n_workers=n_workers,
            max_gpu_batch=max_gpu_batch,
            batch_timeout_ms=batch_timeout_ms,
        )

    def __repr__(self) -> str:
        return (
            f"Geney(splicing_engine={self.splicing_engine!r}, "
            f"organism={self.organism!r})"
        )
