# geney/folding/pipeline.py — High-level pipeline functions for mRNA folding analysis.
"""On-the-fly single-mutation scoring and batch query against precomputed tables."""

from __future__ import annotations

import os

import numpy as np
import pandas as pd

from .delta_mfe import compute_delta_mfe, compute_folding_profile
from .query import (
    DeltaTableIndex,
    query_mutation,
    OK,
)
from .reference import normalize_transcript_id, complement


def _central_window_dict(profile) -> dict:
    """Pick the window closest to centered on the mutation and flatten it."""
    if not profile.windows:
        return {}
    best = min(profile.windows, key=lambda w: abs(w.mutation_offset - len(w.ref_seq) // 2))
    return {
        "ref_window": best.ref_seq,
        "var_window": best.var_seq,
        "ref_structure": best.ref_structure,
        "var_structure": best.var_structure,
        "mfe_ref": best.ref_mfe,
        "mfe_var": best.var_mfe,
        "mutation_offset": best.mutation_offset,
    }

__all__ = [
    "set_folding_data_dir",
    "get_folding_data_dir",
    "folding_pipeline",
    "folding_query_batch",
]


# ---------------------------------------------------------------------------
# Global data-directory configuration (follows the TITER pattern)
# ---------------------------------------------------------------------------

_FOLDING_DATA_DIR: str | None = None


def set_folding_data_dir(path: str) -> None:
    """Set base directory for precomputed folding data.

    Expected layout inside *path*::

        step1/          # per-gene conserved position files
        delta_table/    # precomputed delta-MFE parquet files
        cds_sequences.txt.gz
        cds_genome_map.pickle
    """
    global _FOLDING_DATA_DIR
    _FOLDING_DATA_DIR = path


def get_folding_data_dir() -> str | None:
    """Return the current folding data directory (env ``FOLDING_DATA_DIR`` as fallback)."""
    global _FOLDING_DATA_DIR
    if _FOLDING_DATA_DIR is None:
        _FOLDING_DATA_DIR = os.environ.get("FOLDING_DATA_DIR")
    return _FOLDING_DATA_DIR


# ---------------------------------------------------------------------------
# On-the-fly single-mutation analysis (no precomputed data required)
# ---------------------------------------------------------------------------

def folding_pipeline(
    mut_id: str,
    transcript_id: str | None = None,
    organism: str = "hg38",
    window_size: int = 39,
) -> dict:
    """Compute delta-MFE for a single mutation on-the-fly.

    Loads gene via *seqmat*, extracts the mature mRNA, applies the mutation,
    and computes the mean delta-MFE over overlapping windows.

    **No p-value** is produced in this mode — that requires a precomputed delta
    table (see :func:`folding_query_batch`).

    ViennaRNA must be installed (``pip install ViennaRNA``).

    Args:
        mut_id: Mutation ID (``GENE:CHR:POS:REF:ALT``).
        transcript_id: Optional Ensembl transcript ID.
        organism: Genome build (default ``"hg38"``).
        window_size: Local folding window (default 39 nt).

    Returns:
        ``dict`` with keys ``status``, ``delta_mfe``, ``mut_id``, ``gene``,
        ``transcript_id``, ``mutations``.
    """
    from ..variants import MutationalEvent
    from ..utils import select_transcript

    # --- validate ViennaRNA ---
    from .vienna import fold_available

    from ..results import FoldingResult

    if not fold_available():
        return FoldingResult(status="error", reason="ViennaRNA not installed (pip install ViennaRNA)",
                             mut_id=mut_id)

    # --- parse mutation ---
    m = MutationalEvent(mut_id)
    if not m.compatible():
        raise ValueError("Mutations in event are incompatible")

    # --- load gene ---
    try:
        from seqmat import Gene

        gene = Gene.from_file(m.gene, organism=organism)
    except Exception as e:
        return FoldingResult(status="error", reason=f"gene_load_failed: {e}", mut_id=mut_id)

    central_pos = m.central_position
    selected = select_transcript(gene, central_pos, transcript_id)
    if selected is None:
        return FoldingResult(status="error", reason="no_transcript", mut_id=mut_id)

    # --- build mature mRNA ---
    reference_transcript = selected.generate_mature_mrna()

    mrna = reference_transcript.mature_mrna
    if mrna is None or not hasattr(mrna, "seq") or not hasattr(mrna, "index"):
        return FoldingResult(status="error", reason="no_mature_mrna", mut_id=mut_id)

    ref_seq = mrna.seq.upper()
    idx_list = mrna.index.tolist() if hasattr(mrna.index, "tolist") else list(mrna.index)
    pos_to_mrna = {gpos: i for i, gpos in enumerate(idx_list)}

    # --- score each sub-mutation ---
    per_mut: list[dict] = []
    for pos, ref, alt in m.mutation_args():
        mrna_pos = pos_to_mrna.get(pos)
        if mrna_pos is None:
            per_mut.append(
                {"genomic_position": pos, "ref": ref, "alt": alt, "status": "not_in_mrna", "delta_mfe": None}
            )
            continue

        actual_ref = ref_seq[mrna_pos]
        coding_ref = ref.upper()
        coding_alt = alt.upper()

        # Handle strand: if transcript is reverse-complement, seqmat already
        # stores the mRNA on the coding strand, but the mutation alleles may
        # be reported on the + strand.
        if actual_ref != coding_ref:
            coding_ref = complement(coding_ref)
            coding_alt = complement(coding_alt)
            if actual_ref != coding_ref:
                per_mut.append(
                    {
                        "genomic_position": pos,
                        "ref": ref,
                        "alt": alt,
                        "status": "ref_mismatch",
                        "delta_mfe": None,
                        "expected_ref": actual_ref,
                    }
                )
                continue

        var_seq = ref_seq[:mrna_pos] + coding_alt + ref_seq[mrna_pos + 1 :]

        try:
            profile = compute_folding_profile(ref_seq, var_seq, mrna_pos, window_size)
            entry = {
                "genomic_position": pos,
                "mrna_position": mrna_pos,
                "ref": ref,
                "alt": alt,
                "coding_alt": coding_alt,
                "status": "ok",
                "delta_mfe": profile.delta_mfe,
                "score": profile.score,
            }
            entry.update(_central_window_dict(profile))
            per_mut.append(entry)
        except Exception as e:
            per_mut.append(
                {"genomic_position": pos, "ref": ref, "alt": alt, "status": "fold_error", "delta_mfe": None, "error": str(e)}
            )

    from ..results import FoldingResult, FoldingWindow

    # Build FoldingWindow from the best entry
    central_window = None
    for entry in per_mut:
        if entry.get("status") == "ok" and entry.get("ref_window"):
            central_window = FoldingWindow(
                ref_sequence=entry["ref_window"],
                var_sequence=entry["var_window"],
                ref_structure=entry["ref_structure"],
                var_structure=entry["var_structure"],
                ref_mfe=entry["mfe_ref"],
                var_mfe=entry["mfe_var"],
                delta_mfe=entry["mfe_var"] - entry["mfe_ref"],
                mutation_offset=entry["mutation_offset"],
            )
            break

    first_delta = per_mut[0]["delta_mfe"] if per_mut else None
    first_score = next((e.get("score", 0.0) for e in per_mut if e.get("status") == "ok"), 0.0)

    return FoldingResult(
        status="completed",
        score=first_score or 0.0,
        delta_mfe=first_delta,
        mutations=per_mut,
        central_window=central_window,
        mut_id=mut_id,
        gene=m.gene,
        transcript_id=getattr(reference_transcript, "transcript_id", None),
        window_size=window_size,
    )


# ---------------------------------------------------------------------------
# Batch query with precomputed data (fast — no Vienna calls)
# ---------------------------------------------------------------------------

def folding_query_batch(
    mutations_df: pd.DataFrame,
    delta_table: pd.DataFrame | DeltaTableIndex,
    step1_data: pd.DataFrame | None = None,
    cds_genome_map: dict | None = None,
    cds_sequences: dict[str, str] | None = None,
    require_conserved: bool = True,
    fdr_threshold: float = 0.05,
) -> pd.DataFrame:
    """Score a batch of mutations against precomputed delta-MFE data.

    **No ViennaRNA calls at query time.**  All delta-MFE values are looked up
    from *delta_table*; p-values are computed against the empirical null.

    Args:
        mutations_df: Mutation table with columns ``Start_Position``, ``ENSG``,
            ``ENST``, ``Reference_Allele``, ``Alternative_Allele``, ``Strand``.
        delta_table: Precomputed delta-MFE DataFrame (or a pre-built
            :class:`~.query.DeltaTableIndex`).
        step1_data: Optional Step 1 conserved-position DataFrame for filtering.
        cds_genome_map: Optional CDS-to-genome coordinate map (see
            :func:`~.reference.load_cds_genome_map`).  Only needed when
            ``mutations_df`` lacks a ``cds_position_1based`` column.
        cds_sequences: Optional ``{transcript_id: cds_seq}`` for ref-base
            validation.
        require_conserved: If *True* and *step1_data* provided, skip
            non-conserved positions.
        fdr_threshold: FDR threshold for the conserved position check.

    Returns:
        DataFrame sorted by: scored rows first, then ``p_empirical_two_sided``
        ascending, then ``delta_mfe_true`` descending.
    """
    from .conserved import ConservedPositionIndex
    from .reference import genome_to_cds_position

    # Build indexes
    if not isinstance(delta_table, DeltaTableIndex):
        dt_index = DeltaTableIndex(delta_table)
    else:
        dt_index = delta_table

    cons_index = None
    if step1_data is not None and not step1_data.empty:
        cons_index = ConservedPositionIndex(step1_data, fdr_threshold)

    _COMP = {"A": "T", "T": "A", "C": "G", "G": "C"}

    results: list[dict] = []
    for row_id, row in mutations_df.iterrows():
        enst = str(row.get("ENST", ""))
        ensg = str(row.get("ENSG", ""))
        genome_pos = int(row.get("Start_Position", 0))
        ref_allele = str(row.get("Reference_Allele", "")).upper()
        alt_allele = str(row.get("Alternative_Allele", "")).upper()
        strand = str(row.get("Strand", "+"))

        # --- determine CDS position ---
        cds_pos_1 = row.get("cds_position_1based")
        if cds_pos_1 is None or (isinstance(cds_pos_1, float) and np.isnan(cds_pos_1)):
            if cds_genome_map is not None:
                cds_pos_0 = genome_to_cds_position(cds_genome_map, enst, genome_pos)
                cds_pos_1 = (cds_pos_0 + 1) if cds_pos_0 is not None else None
            else:
                cds_pos_1 = None

        if cds_pos_1 is None:
            results.append(_error_row(row_id, genome_pos, ensg, enst, ref_allele, alt_allele, strand, "ERR_POSITION_NOT_IN_TRANSCRIPT"))
            continue

        cds_pos_1 = int(cds_pos_1)

        # --- resolve coding-strand allele ---
        coding_alt = alt_allele
        if strand == "-":
            coding_alt = _COMP.get(alt_allele, alt_allele)

        # Validate ref against CDS if available
        if cds_sequences is not None:
            tid_norm = normalize_transcript_id(enst)
            cds_seq = cds_sequences.get(tid_norm)
            if cds_seq is not None:
                cds_pos_0 = cds_pos_1 - 1
                if 0 <= cds_pos_0 < len(cds_seq):
                    expected_ref = cds_seq[cds_pos_0].upper()
                    coding_ref = ref_allele if strand != "-" else _COMP.get(ref_allele, ref_allele)
                    if expected_ref != coding_ref:
                        results.append(_error_row(row_id, genome_pos, ensg, enst, ref_allele, alt_allele, strand, "ERR_REF_MISMATCH_CDS"))
                        continue

        # --- query ---
        q = query_mutation(
            delta_table=dt_index,
            transcript_id=enst,
            cds_position_1based=cds_pos_1,
            alt_base=coding_alt,
            conserved_index=cons_index,
            require_conserved=require_conserved,
        )

        results.append(
            {
                "row_id": row_id,
                "Start_Position": genome_pos,
                "ENSG": ensg,
                "ENST": enst,
                "Reference_Allele": ref_allele,
                "Alternative_Allele": alt_allele,
                "Strand": strand,
                "cds_position_1based": cds_pos_1,
                "delta_mfe_true": q["delta_mfe_true"],
                "p_empirical_two_sided": q["p_empirical_two_sided"],
                "status_code": q["status_code"],
                "null_set_size": q["null_set_size"],
                "conserved_state": q.get("conserved_state"),
                "functional_class": q.get("functional_class"),
            }
        )

    result_df = pd.DataFrame(results)

    # Sort: scored first → p ascending → |delta| descending → row_id
    if not result_df.empty and "status_code" in result_df.columns:
        result_df["_scored"] = (result_df["status_code"] == OK).astype(int)
        result_df = (
            result_df.sort_values(
                by=["_scored", "p_empirical_two_sided", "delta_mfe_true", "row_id"],
                ascending=[False, True, False, True],
            )
            .drop(columns=["_scored"])
            .reset_index(drop=True)
        )

    return result_df


def _error_row(row_id, genome_pos, ensg, enst, ref, alt, strand, status_code):
    return {
        "row_id": row_id,
        "Start_Position": genome_pos,
        "ENSG": ensg,
        "ENST": enst,
        "Reference_Allele": ref,
        "Alternative_Allele": alt,
        "Strand": strand,
        "cds_position_1based": None,
        "delta_mfe_true": None,
        "p_empirical_two_sided": None,
        "status_code": status_code,
        "null_set_size": 0,
        "conserved_state": None,
        "functional_class": None,
    }
