# geney/screen_alphagenome.py — AlphaGenome-based fast variant screening.
"""
Fast variant screening using AlphaGenome (JAX).

AlphaGenome predicts splicing, expression, and chromatin effects in a single
forward pass, making it ideal for screening millions of variants on GPU.

Two screening modes:

1. **Embedding delta** (fastest) — Compare mean-pooled embeddings between
   ref and variant sequences.  One scalar capturing all effects at once.

2. **Track delta** (more detailed) — Compare specific output tracks
   (splice sites, RNA-seq) between ref and variant to get per-subsystem deltas.

Requires: alphagenome-research, JAX, NVIDIA GPU (T4+).
"""
from __future__ import annotations

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

import numpy as np

from .variants import MutationalEvent

logger = logging.getLogger(__name__)

__all__ = [
    "AlphaGenomeScreenResult",
    "screen_mutation_alphagenome",
    "screen_mutations_alphagenome",
]

# Lazy-loaded wrapper instance
_ag_wrapper = None


def _get_wrapper(model_version: str = "all_folds", source: str = "huggingface"):
    """Lazy-load AlphaGenome wrapper (singleton)."""
    global _ag_wrapper
    if _ag_wrapper is not None:
        return _ag_wrapper

    from genebeddings.wrappers.alphagenome_wrapper import AlphaGenomeWrapper
    _ag_wrapper = AlphaGenomeWrapper(model_version=model_version, source=source)
    logger.info("AlphaGenome loaded (%s)", model_version)
    return _ag_wrapper


def set_alphagenome_model(model_version: str = "all_folds", source: str = "huggingface"):
    """Pre-configure the AlphaGenome model (call before screening)."""
    global _ag_wrapper
    _ag_wrapper = None
    _get_wrapper(model_version, source)


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

@dataclass(slots=True)
class AlphaGenomeScreenResult:
    """Screening result from AlphaGenome — unified model, single forward pass."""
    mut_id: str
    gene: str
    embedding_delta: float | None   # L2 distance between ref/var embeddings
    splicing_delta: float | None    # max abs delta in splice site tracks
    expression_delta: float | None  # max abs delta in RNA-seq/CAGE tracks
    chromatin_delta: float | None   # max abs delta in ATAC tracks
    error: str | None = None

    @property
    def max_delta(self) -> float:
        vals = [v for v in (self.embedding_delta, self.splicing_delta,
                            self.expression_delta, self.chromatin_delta)
                if v is not None]
        return max(abs(v) for v in vals) if vals else 0.0


# ============================================================================
# Sequence extraction
# ============================================================================

def _extract_ref_var_sequences(
    chrom: str,
    pos: int,
    mutations: list[tuple[int, str, str]],
    organism: str = "hg38",
    seq_len: int = 131_072,
) -> tuple[str, str]:
    """Extract ref and variant genomic sequences centered on pos.

    Uses seqmat for sequence extraction. Returns sequences of exactly
    ``seq_len`` bp.
    """
    from seqmat import SeqMat

    if not chrom.startswith("chr"):
        chrom = f"chr{chrom}"

    half = seq_len // 2
    start = max(1, pos - half)
    end = start + seq_len - 1

    wt_sm = SeqMat.from_fasta(genome=organism, chrom=chrom, start=start, end=end)
    wt_seq = wt_sm.seq

    if len(wt_seq) < seq_len:
        wt_seq = wt_seq + ("N" * (seq_len - len(wt_seq)))

    mut_sm = wt_sm.clone()
    mut_sm.apply_mutations(mutations)
    mut_seq = mut_sm.seq

    # Fix length for indels
    len_diff = len(mut_seq) - seq_len
    if len_diff > 0:
        trim = len_diff // 2
        mut_seq = mut_seq[trim:trim + seq_len]
    elif len_diff < 0:
        mut_seq = mut_seq + ("N" * (-len_diff))

    return wt_seq, mut_seq


# ============================================================================
# Embedding-based screen (fastest)
# ============================================================================

def _screen_embedding_delta(
    wrapper,
    ref_seq: str,
    var_seq: str,
    resolution: int = 1,
) -> float:
    """L2 distance between mean-pooled ref and variant embeddings."""
    ref_emb = wrapper.embed(ref_seq, pool="mean", resolution=resolution)
    var_emb = wrapper.embed(var_seq, pool="mean", resolution=resolution)
    return float(np.linalg.norm(var_emb - ref_emb))


def _screen_track_deltas(
    wrapper,
    ref_seq: str,
    var_seq: str,
) -> dict[str, float | None]:
    """Compute per-track-type max absolute deltas from AlphaGenome predictions.

    Runs predict_tracks on ref and variant, then for each available track type
    computes max |var - ref| across all positions and sub-tracks.

    Returns dict with keys: splicing_delta, expression_delta, chromatin_delta.
    """
    try:
        from alphagenome.models.dna_output import OutputType
    except ImportError:
        from alphagenome_research.model.dna_output import OutputType

    # Request only the tracks we care about for screening
    requested = [
        OutputType.SPLICE_SITES,
        OutputType.RNA_SEQ,
        OutputType.CAGE,
        OutputType.ATAC,
    ]

    ref_output = wrapper.predict_tracks(ref_seq, requested_outputs=requested)
    var_output = wrapper.predict_tracks(var_seq, requested_outputs=requested)

    result = {
        "splicing_delta": None,
        "expression_delta": None,
        "chromatin_delta": None,
    }

    # Splice sites delta
    ref_splice = ref_output.splice_sites
    var_splice = var_output.splice_sites
    if ref_splice is not None and var_splice is not None:
        ref_vals = np.asarray(ref_splice.values)
        var_vals = np.asarray(var_splice.values)
        n = min(ref_vals.shape[0], var_vals.shape[0])
        result["splicing_delta"] = float(np.max(np.abs(var_vals[:n] - ref_vals[:n])))

    # Expression delta (RNA-seq + CAGE combined)
    exp_deltas = []
    for attr in ("rna_seq", "cage"):
        ref_track = getattr(ref_output, attr, None)
        var_track = getattr(var_output, attr, None)
        if ref_track is not None and var_track is not None:
            ref_vals = np.asarray(ref_track.values)
            var_vals = np.asarray(var_track.values)
            n = min(ref_vals.shape[0], var_vals.shape[0])
            exp_deltas.append(float(np.max(np.abs(var_vals[:n] - ref_vals[:n]))))
    if exp_deltas:
        result["expression_delta"] = max(exp_deltas)

    # Chromatin delta (ATAC)
    ref_atac = ref_output.atac
    var_atac = var_output.atac
    if ref_atac is not None and var_atac is not None:
        ref_vals = np.asarray(ref_atac.values)
        var_vals = np.asarray(var_atac.values)
        n = min(ref_vals.shape[0], var_vals.shape[0])
        result["chromatin_delta"] = float(np.max(np.abs(var_vals[:n] - ref_vals[:n])))

    return result


# ============================================================================
# Single mutation screen
# ============================================================================

def screen_mutation_alphagenome(
    mut_id: str,
    organism: str = "hg38",
    seq_len: int = 131_072,
    resolution: int = 1,
    mode: str = "both",
    model_version: str = "all_folds",
    source: str = "huggingface",
) -> AlphaGenomeScreenResult:
    """Screen a single mutation using AlphaGenome.

    Args:
        mut_id: Mutation ID (GENE:CHR:POS:REF:ALT).
        organism: Genome build.
        seq_len: Sequence length to extract (default 131,072 = 128k bp).
        resolution: Embedding resolution (1 or 128 bp).
        mode: ``"embedding"`` (fastest), ``"tracks"`` (decomposed deltas),
            or ``"both"`` (default).
        model_version: AlphaGenome model version.
        source: Model source ('huggingface' or 'kaggle').
    """
    try:
        m = MutationalEvent(mut_id)
        if not m.compatible():
            return AlphaGenomeScreenResult(mut_id, m.gene, None, None, None, None,
                                           error="incompatible_mutations")

        wrapper = _get_wrapper(model_version, source)

        ref_seq, var_seq = _extract_ref_var_sequences(
            str(m.chrom), m.central_position, m.mutation_args(),
            organism, seq_len,
        )

        emb_delta = None
        splicing_delta = None
        expression_delta = None
        chromatin_delta = None

        if mode in ("embedding", "both"):
            emb_delta = _screen_embedding_delta(wrapper, ref_seq, var_seq, resolution)

        if mode in ("tracks", "both"):
            track_deltas = _screen_track_deltas(wrapper, ref_seq, var_seq)
            splicing_delta = track_deltas["splicing_delta"]
            expression_delta = track_deltas["expression_delta"]
            chromatin_delta = track_deltas["chromatin_delta"]

        return AlphaGenomeScreenResult(
            mut_id=mut_id,
            gene=m.gene,
            embedding_delta=emb_delta,
            splicing_delta=splicing_delta,
            expression_delta=expression_delta,
            chromatin_delta=chromatin_delta,
        )

    except Exception as e:
        gene = mut_id.split(":")[0] if ":" in mut_id else ""
        return AlphaGenomeScreenResult(mut_id, gene, None, None, None, None,
                                       error=str(e))


# ============================================================================
# Streaming batch screen
# ============================================================================

def screen_mutations_alphagenome(
    mut_ids: Iterable[str],
    organism: str = "hg38",
    seq_len: int = 131_072,
    resolution: int = 1,
    batch_size: int = 64,
    model_version: str = "all_folds",
    source: str = "huggingface",
    on_result: Callable[[AlphaGenomeScreenResult], None] | None = None,
) -> Generator[AlphaGenomeScreenResult, None, None]:
    """Stream-screen mutations using AlphaGenome embeddings.

    For each mutation: extract ref/var sequences, compute embedding delta.
    Processes mutations in batches for efficient GPU utilization.

    Args:
        mut_ids: Iterable of mutation ID strings.
        organism: Genome build.
        seq_len: Sequence length per mutation (default 128k bp).
        resolution: Embedding resolution (1 or 128 bp).
        batch_size: Mutations to collect before processing.
        model_version: AlphaGenome model version.
        source: Model source.
        on_result: Optional callback per result.

    Yields:
        AlphaGenomeScreenResult per mutation.
    """
    wrapper = _get_wrapper(model_version, source)
    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 _process_batch(batch, wrapper, organism, seq_len,
                                      resolution, on_result)
            batch = []

    if batch:
        yield from _process_batch(batch, wrapper, organism, seq_len,
                                  resolution, on_result)


def _process_batch(
    mut_ids: list[str],
    wrapper,
    organism: str,
    seq_len: int,
    resolution: int,
    on_result: Callable[[AlphaGenomeScreenResult], None] | None,
) -> Generator[AlphaGenomeScreenResult, None, None]:
    """Process a batch — extract sequences, compute embedding deltas."""
    from concurrent.futures import ThreadPoolExecutor

    # Parse mutations and extract sequences in parallel (I/O bound)
    parsed: list[tuple[MutationalEvent, str, str] | None] = []

    def _prepare(mid: str):
        try:
            m = MutationalEvent(mid)
            if not m.compatible():
                return None
            ref_seq, var_seq = _extract_ref_var_sequences(
                str(m.chrom), m.central_position, m.mutation_args(),
                organism, seq_len,
            )
            return (m, ref_seq, var_seq)
        except Exception:
            return None

    with ThreadPoolExecutor(max_workers=min(8, len(mut_ids))) as pool:
        parsed = list(pool.map(_prepare, mut_ids))

    # Compute deltas
    for mid, prep in zip(mut_ids, parsed):
        if prep is None:
            result = AlphaGenomeScreenResult(
                mid, mid.split(":")[0] if ":" in mid else "",
                None, None, None, None, error="preparation_failed",
            )
        else:
            m, ref_seq, var_seq = prep
            try:
                emb_delta = _screen_embedding_delta(wrapper, ref_seq, var_seq, resolution)
                result = AlphaGenomeScreenResult(
                    mut_id=mid, gene=m.gene,
                    embedding_delta=emb_delta,
                    splicing_delta=None, expression_delta=None,
                    chromatin_delta=None,
                )
            except Exception as e:
                result = AlphaGenomeScreenResult(
                    mid, m.gene, None, None, None, None, error=str(e),
                )

        if on_result is not None:
            on_result(result)
        yield result


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

def screen_and_select_alphagenome(
    mut_ids: Iterable[str],
    top_k: int = 100,
    organism: str = "hg38",
    seq_len: int = 131_072,
    resolution: int = 1,
    batch_size: int = 64,
    model_version: str = "all_folds",
    source: str = "huggingface",
) -> list[AlphaGenomeScreenResult]:
    """Screen all mutations and return top-K by embedding delta.

    Uses a min-heap for O(N log K) memory efficiency.
    """
    import heapq

    heap: list[tuple[float, int, AlphaGenomeScreenResult]] = []
    counter = 0

    for result in screen_mutations_alphagenome(
        mut_ids, organism=organism, seq_len=seq_len, resolution=resolution,
        batch_size=batch_size, model_version=model_version, source=source,
    ):
        key = result.max_delta
        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 [r for _, _, r in sorted(heap, reverse=True)]
