# 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="spliceai",              # 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 = "spliceai",
        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_tis: bool = True,
        run_folding: bool = True,
        run_transcription: bool = True,
        enable_borzoi: bool | None = None,
        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,
    ) -> 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_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,
        )

    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 __repr__(self) -> str:
        return (
            f"Geney(splicing_engine={self.splicing_engine!r}, "
            f"organism={self.organism!r})"
        )
