# geney/screen.py — Fast variant screening with minimal-computation deltas.
"""
Lightweight screening pipeline that computes a single scalar delta per
subsystem for each mutation.  Designed to process millions of variants
quickly so that only the top candidates need full analysis.

Each screen function returns a single float:

* **splicing_delta** — max absolute change in splice-site probability
  (ref vs variant) across all positions in the prediction window.
* **tis_delta** — change in canonical start-codon TIS probability.
* **folding_delta** — delta-MFE from a single central window fold.
* **transcription_delta** — max absolute log2 fold-change across GTEx tissues.

The streaming entry point :func:`screen_mutations` accepts an iterable of
mutation IDs and yields :class:`ScreenResult` tuples, internally batching
GPU inference for throughput.
"""
from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Callable, Generator, Iterable

import numpy as np

from seqmat import Gene
from .variants import MutationalEvent
from .utils import (
    select_transcript,
    _apply_mutation_safe,
    _position_in_transcript,
)

logger = logging.getLogger(__name__)

__all__ = [
    "ScreenResult",
    "screen_splicing",
    "screen_tis",
    "screen_folding",
    "screen_transcription",
    "screen_mutation",
    "screen_mutations",
]


# ============================================================================
# Result container
# ============================================================================

@dataclass(slots=True)
class ScreenResult:
    """Lightweight screening result — one float per subsystem."""
    mut_id: str
    gene: str
    splicing_delta: float | None
    tis_delta: float | None
    folding_delta: float | None
    transcription_delta: float | None
    error: str | None = None

    @property
    def max_delta(self) -> float:
        """Largest absolute delta across all subsystems (ignoring None)."""
        vals = [
            v for v in (self.splicing_delta, self.tis_delta,
                        self.folding_delta, self.transcription_delta)
            if v is not None
        ]
        return max(abs(v) for v in vals) if vals else 0.0


# ============================================================================
# Individual screen functions — minimal computation, single scalar output
# ============================================================================

def screen_splicing(
    ref_seq: str,
    var_seq: str,
    engine: str = "spliceai",
) -> float:
    """Max absolute splice-probability delta between ref and variant.

    Runs the splice model once on each sequence and returns the single
    largest absolute change in donor or acceptor probability across all
    positions.  No TranscriptLibrary, no Oncosplice, no DataFrames.
    """
    from .splicing.engines import run_splicing_engine

    ref_don, ref_acc = run_splicing_engine(ref_seq, engine=engine)
    var_don, var_acc = run_splicing_engine(var_seq, engine=engine)

    ref_don, ref_acc = np.asarray(ref_don), np.asarray(ref_acc)
    var_don, var_acc = np.asarray(var_don), np.asarray(var_acc)

    # Truncate to common length (should match, but be safe)
    n = min(len(ref_don), len(var_don))
    don_delta = np.max(np.abs(var_don[:n] - ref_don[:n]))
    acc_delta = np.max(np.abs(var_acc[:n] - ref_acc[:n]))

    return float(max(don_delta, acc_delta))


def screen_splicing_from_probs(
    ref_don: np.ndarray, ref_acc: np.ndarray,
    var_don: np.ndarray, var_acc: np.ndarray,
) -> float:
    """Max absolute splice-probability delta from pre-computed probabilities."""
    n = min(len(ref_don), len(var_don))
    don_delta = np.max(np.abs(var_don[:n] - ref_don[:n]))
    acc_delta = np.max(np.abs(var_acc[:n] - ref_acc[:n]))
    return float(max(don_delta, acc_delta))


def screen_tis(
    ref_mrna: str,
    var_mrna: str,
) -> float:
    """Delta in canonical start-codon TIS probability (variant - reference).

    Scores only the first ATG in each sequence — no full ATG scan, no
    usage fractions, no DataFrame construction.
    """
    from .tis.prediction import predict_context_scores_batch, extract_tis_context

    # Find canonical ATG
    ref_start = ref_mrna.upper().find("ATG")
    var_start = var_mrna.upper().find("ATG")

    if ref_start == -1:
        return 0.0

    contexts = [extract_tis_context(ref_mrna, ref_start)]
    if var_start != -1:
        contexts.append(extract_tis_context(var_mrna, var_start))

    scores = predict_context_scores_batch(contexts)

    ref_score = float(scores[0])
    var_score = float(scores[1]) if var_start != -1 else 0.0

    return var_score - ref_score


def screen_folding(
    ref_seq: str,
    var_seq: str,
    position_0based: int,
    window_size: int = 39,
) -> float:
    """Delta-MFE from a single central window (not all overlapping windows).

    One ViennaRNA fold call per sequence instead of ~window_size calls.
    """
    from .folding.vienna import fold

    n = len(ref_seq)
    if position_0based < 0 or position_0based >= n:
        return 0.0

    # Single central window
    half = window_size // 2
    start = max(0, position_0based - half)
    end = min(n, start + window_size)
    start = max(0, end - window_size)  # adjust if near end

    ref_win = ref_seq[start:end]
    var_win = var_seq[start:end]

    if ref_win == var_win:
        return 0.0

    return fold(var_win) - fold(ref_win)


def screen_transcription(
    chrom: str,
    pos: int,
    mutations: list[tuple[int, str, str]],
    organism: str = "hg38",
) -> float:
    """Max absolute log2 fold-change across GTEx tissues.

    Runs Borzoi once (2-sequence batch) and returns a single scalar.
    """
    from .transcription.prediction import predict_transcription_effect
    result = predict_transcription_effect(chrom, pos, mutations, organism)
    return float(result.get("max_abs_lfc", 0.0))


# ============================================================================
# Single-mutation screen (convenience)
# ============================================================================

def screen_mutation(
    mut_id: str,
    organism: str = "hg38",
    splicing_engine: str = "spliceai",
    run_splicing: bool = True,
    run_tis: bool = True,
    run_folding: bool = True,
    run_transcription: bool = True,
    folding_window_size: int = 39,
) -> ScreenResult:
    """Screen a single mutation across all subsystems."""
    try:
        m = MutationalEvent(mut_id)
        if not m.compatible():
            return ScreenResult(mut_id, m.gene, None, None, None, None,
                                error="incompatible_mutations")
        central_pos = m.central_position

        gene = Gene.from_file(m.gene, organism=organism)
        selected = select_transcript(gene, central_pos)
        if selected is None:
            return ScreenResult(mut_id, m.gene, None, None, None, None,
                                error="no_transcript")

        # Build ref and variant via direct-exon assembly (no full pre-mRNA load)
        ref_transcript = selected.clone().generate_mature_mrna()
        var_transcript = selected.clone().generate_mature_mrna()
        for pos, ref, alt in m.mutation_args():
            try:
                _apply_mutation_safe(var_transcript.mature_mrna, pos, ref, alt)
            except Exception:
                pass

        # --- Splicing screen ---
        splicing_delta = None
        if run_splicing:
            try:
                context = 7500
                ref_for_splicing = selected.clone()
                ref_for_splicing.generate_pre_mrna(
                    region_start=central_pos - context,
                    region_end=central_pos + context,
                )
                var_for_splicing = selected.clone()
                var_for_splicing.generate_pre_mrna(
                    region_start=central_pos - context,
                    region_end=central_pos + context,
                )
                for pos, ref, alt in m.mutation_args():
                    try:
                        _apply_mutation_safe(var_for_splicing.pre_mrna, pos, ref, alt)
                    except Exception:
                        pass
                splicing_delta = screen_splicing(
                    ref_for_splicing.pre_mrna.seq, var_for_splicing.pre_mrna.seq,
                    engine=splicing_engine,
                )
            except Exception as e:
                logger.debug("Splicing screen failed for %s: %s", mut_id, e)

        # --- TIS screen ---
        tis_delta = None
        if run_tis:
            try:
                ref_mrna = ref_transcript.mature_mrna
                var_mrna = var_transcript.mature_mrna
                if ref_mrna is not None and var_mrna is not None:
                    tis_delta = screen_tis(ref_mrna.seq, var_mrna.seq)
            except Exception as e:
                logger.debug("TIS screen failed for %s: %s", mut_id, e)

        # --- Folding screen ---
        folding_delta = None
        if run_folding:
            try:
                from .folding.vienna import fold_available
                from .folding.reference import complement

                if fold_available():
                    ref_m = ref_transcript.mature_mrna
                    var_m = var_transcript.mature_mrna

                    if ref_m is not None and var_m is not None and hasattr(ref_m, 'index'):
                        ref_seq = ref_m.seq.upper()
                        idx_list = ref_m.index.tolist() if hasattr(ref_m.index, 'tolist') else list(ref_m.index)
                        pos_to_mrna = {gpos: i for i, gpos in enumerate(idx_list)}

                        # Use first mutation position that maps to mRNA
                        for pos, ref_allele, alt_allele in m.mutation_args():
                            mrna_pos = pos_to_mrna.get(pos)
                            if mrna_pos is not None:
                                actual_ref = ref_seq[mrna_pos]
                                coding_alt = alt_allele.upper()

                                if actual_ref != ref_allele.upper():
                                    coding_alt = complement(alt_allele.upper())

                                var_seq = ref_seq[:mrna_pos] + coding_alt + ref_seq[mrna_pos + 1:]
                                folding_delta = screen_folding(
                                    ref_seq, var_seq, mrna_pos, folding_window_size
                                )
                                break
            except Exception as e:
                logger.debug("Folding screen failed for %s: %s", mut_id, e)

        # --- Transcription screen ---
        transcription_delta = None
        if run_transcription:
            try:
                transcription_delta = screen_transcription(
                    str(m.chrom), central_pos, m.mutation_args(), organism
                )
            except Exception as e:
                logger.debug("Transcription screen failed for %s: %s", mut_id, e)

        return ScreenResult(
            mut_id=mut_id,
            gene=m.gene,
            splicing_delta=splicing_delta,
            tis_delta=tis_delta,
            folding_delta=folding_delta,
            transcription_delta=transcription_delta,
        )

    except Exception as e:
        gene_name = ""
        try:
            gene_name = mut_id.split(":")[0]
        except Exception:
            pass
        return ScreenResult(mut_id, gene_name, None, None, None, None,
                            error=str(e))


# ============================================================================
# Streaming batch screener
# ============================================================================

def screen_mutations(
    mut_ids: Iterable[str],
    organism: str = "hg38",
    splicing_engine: str = "spliceai",
    batch_size: int = 64,
    run_splicing: bool = True,
    run_tis: bool = True,
    run_folding: bool = True,
    run_transcription: bool = True,
    folding_window_size: int = 39,
    on_result: Callable[[ScreenResult], None] | None = None,
) -> Generator[ScreenResult, None, None]:
    """Stream-screen mutations with batched GPU inference.

    Accepts any iterable of mutation IDs (file handle, generator, list).
    Internally collects ``batch_size`` mutations, batches GPU inference
    (splicing + TIS), then yields one :class:`ScreenResult` per mutation.

    Args:
        mut_ids: Iterable of mutation ID strings.
        organism: Genome build.
        splicing_engine: ``"spliceai"`` or ``"pangolin"``.
        batch_size: Mutations per GPU batch.
        run_splicing: Enable splicing screening.
        run_tis: Enable TIS screening.
        run_folding: Enable folding screening.
        run_transcription: Enable transcription screening.
        folding_window_size: Window size for folding screen.
        on_result: Optional callback invoked with each ScreenResult
            (e.g. to write to disk or database).

    Yields:
        :class:`ScreenResult` per mutation.
    """
    batch: list[str] = []

    for mid in mut_ids:
        mid = mid.strip()
        if not mid:
            continue
        batch.append(mid)
        if len(batch) >= batch_size:
            yield from _screen_batch(
                batch, organism, splicing_engine,
                run_splicing, run_tis, run_folding, run_transcription,
                folding_window_size, on_result,
            )
            batch = []

    if batch:
        yield from _screen_batch(
            batch, organism, splicing_engine,
            run_splicing, run_tis, run_folding, run_transcription,
            folding_window_size, on_result,
        )


def _screen_batch(
    mut_ids: list[str],
    organism: str,
    splicing_engine: str,
    run_splicing: bool,
    run_tis: bool,
    run_folding: bool,
    run_transcription: bool,
    folding_window_size: int,
    on_result: Callable[[ScreenResult], None] | None,
) -> Generator[ScreenResult, None, None]:
    """Process a batch of mutations with batched GPU inference where possible.

    Batches splicing and TIS inference across the batch. Folding and
    transcription are per-mutation (CPU-bound and heavy-GPU respectively).
    """
    # Phase 1: Parse and prepare all mutations
    prepared: list[dict | None] = []
    for mid in mut_ids:
        try:
            m = MutationalEvent(mid)
            if not m.compatible():
                prepared.append(None)
                continue

            central_pos = m.central_position
            gene = Gene.from_file(m.gene, organism=organism)
            selected = select_transcript(gene, central_pos)
            if selected is None:
                prepared.append(None)
                continue

            ref_transcript = selected.clone().generate_mature_mrna()
            var_transcript = selected.clone().generate_mature_mrna()
            for pos, ref, alt in m.mutation_args():
                try:
                    _apply_mutation_safe(var_transcript.mature_mrna, pos, ref, alt)
                except Exception:
                    pass

            prepared.append({
                "event": m,
                "selected": selected,
                "ref_transcript": ref_transcript,
                "var_transcript": var_transcript,
                "central_pos": central_pos,
            })
        except Exception:
            prepared.append(None)

    # Phase 2: Batch splicing — collect all ref/var sequences, run model once
    splice_results: list[float | None] = [None] * len(mut_ids)
    if run_splicing:
        splice_results = _batch_screen_splicing(prepared, splicing_engine)

    # Phase 3: Batch TIS — collect all ref/var mRNA pairs
    tis_results: list[float | None] = [None] * len(mut_ids)
    if run_tis:
        tis_results = _batch_screen_tis(prepared)

    # Phase 4: Folding (CPU, per-mutation — ViennaRNA is fast for single windows)
    folding_results: list[float | None] = [None] * len(mut_ids)
    if run_folding:
        folding_results = _batch_screen_folding(prepared, folding_window_size)

    # Phase 5: Transcription (heavy GPU — per-mutation for now)
    tx_results: list[float | None] = [None] * len(mut_ids)
    if run_transcription:
        tx_results = _batch_screen_transcription(prepared, organism)

    # Yield results
    for i, mid in enumerate(mut_ids):
        p = prepared[i]
        result = ScreenResult(
            mut_id=mid,
            gene=p["event"].gene if p else mid.split(":")[0] if ":" in mid else "",
            splicing_delta=splice_results[i],
            tis_delta=tis_results[i],
            folding_delta=folding_results[i],
            transcription_delta=tx_results[i],
            error=None if p else "preparation_failed",
        )
        if on_result is not None:
            on_result(result)
        yield result


# ============================================================================
# Batched screen helpers
# ============================================================================

def _batch_screen_splicing(
    prepared: list[dict | None],
    engine: str,
) -> list[float | None]:
    """Batch splicing screen: run all ref+var sequences through the model once."""
    from .splicing.engines import batch_run_splicing_engine

    # Collect sequences and track indices
    seqs: list[str] = []
    seq_map: list[tuple[int, str]] = []  # (prepared_index, "ref"|"var")
    context = 7500

    for i, p in enumerate(prepared):
        if p is None:
            continue
        central_pos = p["central_pos"]
        selected = p.get("selected") or p["ref_transcript"]
        try:
            ref_for_splicing = selected.clone()
            ref_for_splicing.generate_pre_mrna(
                region_start=central_pos - context,
                region_end=central_pos + context,
            )
            var_for_splicing = selected.clone()
            var_for_splicing.generate_pre_mrna(
                region_start=central_pos - context,
                region_end=central_pos + context,
            )
            for pos, ref, alt in p["event"].mutation_args():
                try:
                    _apply_mutation_safe(var_for_splicing.pre_mrna, pos, ref, alt)
                except Exception:
                    pass
            seqs.append(ref_for_splicing.pre_mrna.seq)
            seq_map.append((i, "ref"))
            seqs.append(var_for_splicing.pre_mrna.seq)
            seq_map.append((i, "var"))
        except Exception:
            pass

    if not seqs:
        return [None] * len(prepared)

    # Batched model inference
    try:
        all_probs = batch_run_splicing_engine(seqs, engine=engine)
    except Exception as e:
        logger.warning("Batch splicing screen failed: %s", e)
        return [None] * len(prepared)

    # Scatter results back
    results: list[float | None] = [None] * len(prepared)
    prob_by_idx: dict[int, dict] = {}

    for (idx, role), (don_probs, acc_probs) in zip(seq_map, all_probs):
        if idx not in prob_by_idx:
            prob_by_idx[idx] = {}
        prob_by_idx[idx][role] = (np.asarray(don_probs), np.asarray(acc_probs))

    for idx, probs in prob_by_idx.items():
        if "ref" in probs and "var" in probs:
            results[idx] = screen_splicing_from_probs(
                probs["ref"][0], probs["ref"][1],
                probs["var"][0], probs["var"][1],
            )

    return results


def _batch_screen_tis(
    prepared: list[dict | None],
) -> list[float | None]:
    """Batch TIS screen: collect all canonical ATG contexts, predict in one pass."""
    from .tis.prediction import predict_context_scores_batch, extract_tis_context

    contexts: list[str] = []
    context_map: list[tuple[int, str]] = []  # (prepared_index, "ref"|"var")

    for i, p in enumerate(prepared):
        if p is None:
            continue
        try:
            ref_t = p["ref_transcript"]
            var_t = p["var_transcript"]

            ref_mrna = ref_t.mature_mrna
            var_mrna = var_t.mature_mrna

            if ref_mrna is None or var_mrna is None:
                continue

            ref_seq = ref_mrna.seq
            var_seq = var_mrna.seq

            ref_start = ref_seq.upper().find("ATG")
            var_start = var_seq.upper().find("ATG")

            if ref_start == -1:
                continue

            contexts.append(extract_tis_context(ref_seq, ref_start))
            context_map.append((i, "ref"))

            if var_start != -1:
                contexts.append(extract_tis_context(var_seq, var_start))
                context_map.append((i, "var"))

        except Exception:
            pass

    if not contexts:
        return [None] * len(prepared)

    # Single mega-batch prediction
    try:
        all_scores = predict_context_scores_batch(contexts)
    except Exception as e:
        logger.warning("Batch TIS screen failed: %s", e)
        return [None] * len(prepared)

    # Scatter results
    results: list[float | None] = [None] * len(prepared)
    scores_by_idx: dict[int, dict[str, float]] = {}

    for j, (idx, role) in enumerate(context_map):
        if idx not in scores_by_idx:
            scores_by_idx[idx] = {}
        scores_by_idx[idx][role] = float(all_scores[j])

    for idx, scores in scores_by_idx.items():
        ref_score = scores.get("ref", 0.0)
        var_score = scores.get("var", 0.0)
        results[idx] = var_score - ref_score

    return results


def _batch_screen_folding(
    prepared: list[dict | None],
    window_size: int,
) -> list[float | None]:
    """Per-mutation folding screen using single central window."""
    results: list[float | None] = [None] * len(prepared)

    try:
        from .folding.vienna import fold_available
        if not fold_available():
            return results
        from .folding.reference import complement
    except ImportError:
        return results

    for i, p in enumerate(prepared):
        if p is None:
            continue
        try:
            ref_t = p["ref_transcript"]
            ref_m = ref_t.mature_mrna
            if ref_m is None or not hasattr(ref_m, 'index'):
                continue

            ref_seq = ref_m.seq.upper()
            idx_list = ref_m.index.tolist() if hasattr(ref_m.index, 'tolist') else list(ref_m.index)
            pos_to_mrna = {gpos: j for j, gpos in enumerate(idx_list)}

            m = p["event"]
            for pos, ref_allele, alt_allele in m.mutation_args():
                mrna_pos = pos_to_mrna.get(pos)
                if mrna_pos is not None:
                    actual_ref = ref_seq[mrna_pos]
                    coding_alt = alt_allele.upper()
                    if actual_ref != ref_allele.upper():
                        coding_alt = complement(alt_allele.upper())

                    var_seq = ref_seq[:mrna_pos] + coding_alt + ref_seq[mrna_pos + 1:]
                    results[i] = screen_folding(ref_seq, var_seq, mrna_pos, window_size)
                    break
        except Exception:
            pass

    return results


def _batch_screen_transcription(
    prepared: list[dict | None],
    organism: str,
) -> list[float | None]:
    """Batch transcription screen using Borzoi batch inference."""
    from .transcription.model import _load_borzoi
    from .transcription.prediction import (
        _extract_sequences,
        _predict_tracks_batch,
        _build_gtex_mapping,
        _compute_tissue_changes,
    )
    from concurrent.futures import ThreadPoolExecutor

    # Collect specs
    specs: list[dict] = []
    spec_indices: list[int] = []

    for i, p in enumerate(prepared):
        if p is None:
            continue
        try:
            m = p["event"]
            specs.append({
                "chrom": str(m.chrom),
                "pos": p["central_pos"],
                "mutations": m.mutation_args(),
                "organism": organism,
            })
            spec_indices.append(i)
        except Exception:
            pass

    results: list[float | None] = [None] * len(prepared)
    if not specs:
        return results

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

        # Parallel sequence extraction
        def _extract(spec):
            return _extract_sequences(
                spec["chrom"], spec["pos"], spec["mutations"],
                spec.get("organism", "hg38"),
            )

        with ThreadPoolExecutor(max_workers=min(8, len(specs))) as pool:
            seq_pairs = list(pool.map(_extract, specs))

        # Flatten and batch predict
        all_seqs = []
        for wt, mut in seq_pairs:
            all_seqs.append(wt)
            all_seqs.append(mut)

        max_batch = 8
        all_tracks = []
        for j in range(0, len(all_seqs), max_batch):
            chunk = all_seqs[j:j + max_batch]
            all_tracks.extend(
                _predict_tracks_batch(model, chunk, device, dtype, use_autocast)
            )

        # Compute per-mutation max_abs_lfc
        for j, idx in enumerate(spec_indices):
            wt_tracks = all_tracks[j * 2]
            mut_tracks = all_tracks[j * 2 + 1]
            tissue_data = _compute_tissue_changes(wt_tracks, mut_tracks, gtex_map)
            results[idx] = float(tissue_data.get("max_abs_lfc", 0.0))

    except Exception as e:
        logger.warning("Batch transcription screen failed: %s", e)

    return results


# ============================================================================
# Top-K selection helper
# ============================================================================

def screen_and_select(
    mut_ids: Iterable[str],
    top_k: int = 100,
    organism: str = "hg38",
    splicing_engine: str = "spliceai",
    batch_size: int = 64,
    run_splicing: bool = True,
    run_tis: bool = True,
    run_folding: bool = True,
    run_transcription: bool = True,
    sort_by: str = "max_delta",
) -> list[ScreenResult]:
    """Screen all mutations and return the top-K by delta magnitude.

    Args:
        mut_ids: Iterable of mutation ID strings.
        top_k: Number of top candidates to return.
        organism: Genome build.
        splicing_engine: Splice engine to use.
        batch_size: Mutations per GPU batch.
        run_splicing: Enable splicing screening.
        run_tis: Enable TIS screening.
        run_folding: Enable folding screening.
        run_transcription: Enable transcription screening.
        sort_by: Field to sort by — ``"max_delta"`` (default),
            ``"splicing_delta"``, ``"tis_delta"``, ``"folding_delta"``,
            or ``"transcription_delta"``.

    Returns:
        Top-K :class:`ScreenResult` objects sorted by descending delta.
    """
    import heapq

    def _sort_key(r: ScreenResult) -> float:
        if sort_by == "max_delta":
            return r.max_delta
        val = getattr(r, sort_by, None)
        return abs(val) if val is not None else 0.0

    # Use a min-heap of size top_k for memory efficiency
    heap: list[tuple[float, int, ScreenResult]] = []
    counter = 0

    for result in screen_mutations(
        mut_ids, organism=organism, splicing_engine=splicing_engine,
        batch_size=batch_size, run_splicing=run_splicing, run_tis=run_tis,
        run_folding=run_folding, run_transcription=run_transcription,
    ):
        key = _sort_key(result)
        counter += 1
        if len(heap) < top_k:
            heapq.heappush(heap, (key, counter, result))
        elif key > heap[0][0]:
            heapq.heapreplace(heap, (key, counter, result))

    # Return sorted descending
    return [r for _, _, r in sorted(heap, reverse=True)]
