# geney/folding/query.py — Step 2 query: score mutations from precomputed delta-MFE data.
"""Look up delta-MFE, build null set, and compute empirical p-value — no Vienna calls."""

from __future__ import annotations

import numpy as np
import pandas as pd

from .reference import normalize_transcript_id

__all__ = [
    "DeltaTableIndex",
    "query_mutation",
    "compute_empirical_p_value",
    # Status codes
    "OK",
    "ERR_NOT_CONSERVED_FOLDING_POSITION",
    "ERR_TRANSCRIPT_NOT_PRECOMPUTED",
    "ERR_POSITION_NOT_IN_TRANSCRIPT",
    "ERR_REF_MISMATCH_CDS",
    "ERR_INVALID_ALLELE",
    "ERR_NULL_POOL_EMPTY",
    "ERR_INTERNAL",
]


# ---------------------------------------------------------------------------
# Status codes (string enum)
# ---------------------------------------------------------------------------

OK = "OK"
ERR_NOT_CONSERVED_FOLDING_POSITION = "ERR_NOT_CONSERVED_FOLDING_POSITION"
ERR_TRANSCRIPT_NOT_PRECOMPUTED = "ERR_TRANSCRIPT_NOT_PRECOMPUTED"
ERR_POSITION_NOT_IN_TRANSCRIPT = "ERR_POSITION_NOT_IN_TRANSCRIPT"
ERR_REF_MISMATCH_CDS = "ERR_REF_MISMATCH_CDS"
ERR_INVALID_ALLELE = "ERR_INVALID_ALLELE"
ERR_NULL_POOL_EMPTY = "ERR_NULL_POOL_EMPTY"
ERR_INTERNAL = "ERR_INTERNAL"


# ---------------------------------------------------------------------------
# Empirical p-value
# ---------------------------------------------------------------------------

def compute_empirical_p_value(delta_true: float, null_deltas: np.ndarray) -> float:
    """Two-sided empirical p-value.

    ``p = (1 + #{|null| >= |delta_true|}) / (N + 1)``

    Returns ``NaN`` if *null_deltas* is empty.
    """
    n = len(null_deltas)
    if n == 0:
        return float("nan")
    count = int(np.sum(np.abs(null_deltas) >= abs(delta_true)))
    return (1 + count) / (n + 1)


# ---------------------------------------------------------------------------
# DeltaTableIndex — pre-indexed delta table for fast batch lookups
# ---------------------------------------------------------------------------

class DeltaTableIndex:
    """Pre-indexed wrapper around a precomputed delta-MFE DataFrame.

    Building the index normalises transcript IDs once and groups rows for
    O(1) delta-true lookup and O(1) null-set retrieval.
    """

    def __init__(self, delta_table: pd.DataFrame):
        df = delta_table.copy()
        df["norm_tid"] = df["transcript_id"].map(normalize_transcript_id)

        # Fast delta-true lookup: (tid, pos, alt) -> {delta_mfe, functional_class, ref_base}
        self._delta_idx: dict[tuple[str, int, str], dict] = {}
        for row in df.itertuples(index=False):
            key = (row.norm_tid, int(row.cds_position_1based), row.alt_base)
            self._delta_idx[key] = {
                "delta_mfe": float(row.delta_mfe),
                "functional_class": row.functional_class,
                "ref_base": row.ref_base,
            }

        # Null-set groups: (tid, ref, alt, func_class) -> Nx2 ndarray [[pos, delta], ...]
        groups: dict[tuple, list[list]] = {}
        for row in df.itertuples(index=False):
            gkey = (row.norm_tid, row.ref_base, row.alt_base, row.functional_class)
            groups.setdefault(gkey, []).append(
                [int(row.cds_position_1based), float(row.delta_mfe)]
            )
        self._null_groups: dict[tuple, np.ndarray] = {
            k: np.array(v) for k, v in groups.items()
        }

    # ---- lookups ----------------------------------------------------------

    def get_delta_true(
        self, transcript_id: str, cds_position_1based: int, alt_base: str
    ) -> dict | None:
        """Return ``{delta_mfe, functional_class, ref_base}`` or *None*."""
        tid = normalize_transcript_id(transcript_id)
        return self._delta_idx.get((tid, cds_position_1based, alt_base.upper()))

    def get_null_deltas(
        self,
        transcript_id: str,
        ref_base: str,
        alt_base: str,
        functional_class: str,
        exclude_position_1based: int,
    ) -> np.ndarray:
        """Return 1-D array of delta-MFE values for the null set (excluding the true position)."""
        tid = normalize_transcript_id(transcript_id)
        arr = self._null_groups.get(
            (tid, ref_base.upper(), alt_base.upper(), functional_class)
        )
        if arr is None:
            return np.array([], dtype=float)
        mask = arr[:, 0] != exclude_position_1based
        return arr[mask, 1]

    def __repr__(self) -> str:
        return f"DeltaTableIndex({len(self._delta_idx)} entries)"


# ---------------------------------------------------------------------------
# Single-mutation query
# ---------------------------------------------------------------------------

def query_mutation(
    delta_table: pd.DataFrame | DeltaTableIndex,
    transcript_id: str,
    cds_position_1based: int,
    alt_base: str,
    conserved_index=None,
    require_conserved: bool = True,
) -> dict:
    """Score one mutation against a precomputed delta-MFE table.

    Args:
        delta_table: Precomputed delta DataFrame **or** a :class:`DeltaTableIndex`.
        transcript_id: Ensembl transcript ID.
        cds_position_1based: 1-based CDS position of the mutation.
        alt_base: Alternative base on the coding strand (A/C/G/T).
        conserved_index: Optional :class:`~.conserved.ConservedPositionIndex`.
        require_conserved: If *True* and *conserved_index* is given, mutations at
            non-conserved positions get ``ERR_NOT_CONSERVED_FOLDING_POSITION``.

    Returns:
        Dict with keys: ``delta_mfe_true``, ``p_empirical_two_sided``,
        ``status_code``, ``null_set_size``, ``conserved_state``,
        ``functional_class``, ``ref_base``.
    """
    if not isinstance(delta_table, DeltaTableIndex):
        delta_table = DeltaTableIndex(delta_table)

    tid = normalize_transcript_id(transcript_id)
    alt = alt_base.upper()

    result: dict = {
        "transcript_id": tid,
        "cds_position_1based": cds_position_1based,
        "alt_base": alt,
        "delta_mfe_true": None,
        "p_empirical_two_sided": None,
        "status_code": OK,
        "null_set_size": 0,
        "conserved_state": None,
        "functional_class": None,
        "ref_base": None,
    }

    # Step 1 filter
    if conserved_index is not None:
        is_cons, state = conserved_index.lookup(tid, cds_position_1based)
        result["conserved_state"] = state
        if require_conserved and not is_cons:
            result["status_code"] = ERR_NOT_CONSERVED_FOLDING_POSITION
            return result

    # Delta-true lookup
    info = delta_table.get_delta_true(tid, cds_position_1based, alt)
    if info is None:
        result["status_code"] = ERR_TRANSCRIPT_NOT_PRECOMPUTED
        return result

    delta_true = info["delta_mfe"]
    func_class = info["functional_class"]
    ref_base = info["ref_base"]

    result["delta_mfe_true"] = delta_true
    result["functional_class"] = func_class
    result["ref_base"] = ref_base

    # Null set
    null_deltas = delta_table.get_null_deltas(
        tid, ref_base, alt, func_class, cds_position_1based
    )
    result["null_set_size"] = len(null_deltas)

    if len(null_deltas) == 0:
        result["status_code"] = ERR_NULL_POOL_EMPTY
        return result

    # p-value
    result["p_empirical_two_sided"] = compute_empirical_p_value(delta_true, null_deltas)

    return result
