# geney/folding/conserved.py — Step 1 conserved-folding-position loading and lookup.
"""Load Step 1 results and efficiently check whether a CDS position is conserved."""

from __future__ import annotations

from pathlib import Path

import pandas as pd

from .reference import normalize_transcript_id

__all__ = [
    "load_step1_results",
    "is_conserved",
    "ConservedPositionIndex",
]


# ---------------------------------------------------------------------------
# Loading
# ---------------------------------------------------------------------------

def load_step1_results(
    path: str | Path,
    gene_id: str | None = None,
    transcript_id: str | None = None,
) -> pd.DataFrame:
    """Load Step 1 conserved-folding-position results.

    Supported layouts:

    * **Per-gene directory** (``path/`` contains one file per gene):
      ``{gene_id}.csv``, ``.csv.gz``, or ``.parquet``.
    * **Single file**: CSV / CSV.GZ / Parquet with all genes.

    Expected columns (at least one set):

    * ``transcript_id``, ``cds_position_1based``
    * ``conserved_state`` **or** (``significant_low_both`` / ``significant_high_both``
      + ``p_value_fdr_bh_column`` / ``p_value_fdr_bh_vertical``)

    Args:
        path: File or directory with Step 1 output.
        gene_id: When *path* is a directory, load only this gene's file.
        transcript_id: Optional filter (applied after loading).
    """
    path = Path(path)

    if path.is_dir():
        df = _load_from_directory(path, gene_id)
    else:
        df = _read_file(path)
        if df is None:
            return pd.DataFrame()

    if transcript_id and "transcript_id" in df.columns:
        norm = normalize_transcript_id(transcript_id)
        df = df[df["transcript_id"].map(normalize_transcript_id) == norm]

    return df.reset_index(drop=True)


def _load_from_directory(dirpath: Path, gene_id: str | None) -> pd.DataFrame:
    if gene_id:
        norm_gene = normalize_transcript_id(gene_id)  # works for ENSG too
        for ext in (".csv", ".csv.gz", ".parquet", ".tsv", ".tsv.gz"):
            for candidate in (gene_id, norm_gene):
                p = dirpath / f"{candidate}{ext}"
                if p.exists():
                    df = _read_file(p)
                    if df is not None:
                        return df
        return pd.DataFrame()

    # Load all files in the directory
    frames: list[pd.DataFrame] = []
    for p in sorted(dirpath.iterdir()):
        if p.name.startswith("."):
            continue
        if p.suffix in (".csv", ".gz", ".parquet", ".tsv"):
            df = _read_file(p)
            if df is not None and not df.empty:
                frames.append(df)
    return pd.concat(frames, ignore_index=True) if frames else pd.DataFrame()


def _read_file(path: Path) -> pd.DataFrame | None:
    try:
        if path.suffix == ".parquet" or str(path).endswith(".parquet"):
            return pd.read_parquet(path)
        sep = "\t" if ".tsv" in path.name else ","
        return pd.read_csv(path, sep=sep)
    except Exception:
        return None


# ---------------------------------------------------------------------------
# Single-position check (convenience)
# ---------------------------------------------------------------------------

def is_conserved(
    step1_df: pd.DataFrame,
    transcript_id: str,
    cds_position_1based: int,
    fdr_threshold: float = 0.05,
) -> tuple[bool, str]:
    """Check whether a single CDS position is conserved (low or high MFE).

    Returns:
        ``(True, 'conserved_low_mfe')`` or ``(True, 'conserved_high_mfe')``
        or ``(False, 'not_conserved')``.
    """
    tid = normalize_transcript_id(transcript_id)
    pos_col = "cds_position_1based" if "cds_position_1based" in step1_df.columns else "position"

    mask = step1_df[pos_col] == cds_position_1based
    if "transcript_id" in step1_df.columns:
        mask = mask & (step1_df["transcript_id"].map(normalize_transcript_id) == tid)

    rows = step1_df[mask]
    if rows.empty:
        return False, "not_conserved"

    row = rows.iloc[0]
    return _classify_row(row, fdr_threshold)


def _classify_row(row, fdr_threshold: float) -> tuple[bool, str]:
    """Determine conserved state from a single Step 1 row."""
    # Direct conserved_state column (if present)
    if "conserved_state" in row.index:
        state = str(row["conserved_state"])
        if state in ("conserved_low_mfe", "conserved_high_mfe"):
            return True, state
        return False, "not_conserved"

    # Derive from FDR + significance flags
    fdr_col = row.get("p_value_fdr_bh_column", 1.0)
    fdr_vert = row.get("p_value_fdr_bh_vertical", 1.0)

    try:
        if float(fdr_col) >= fdr_threshold or float(fdr_vert) >= fdr_threshold:
            return False, "not_conserved"
    except (TypeError, ValueError):
        return False, "not_conserved"

    if row.get("significant_low_both", False):
        return True, "conserved_low_mfe"
    if row.get("significant_high_both", False):
        return True, "conserved_high_mfe"

    return False, "not_conserved"


# ---------------------------------------------------------------------------
# Fast pre-built index for batch queries
# ---------------------------------------------------------------------------

class ConservedPositionIndex:
    """Pre-built lookup index for conserved folding positions.

    Build once, query many times — avoids repeated DataFrame filtering.
    """

    def __init__(self, step1_df: pd.DataFrame, fdr_threshold: float = 0.05):
        self._index: dict[tuple[str, int], str] = {}
        self._fdr_threshold = fdr_threshold
        self._build(step1_df)

    # ---- construction -----------------------------------------------------

    def _build(self, df: pd.DataFrame) -> None:
        pos_col = "cds_position_1based" if "cds_position_1based" in df.columns else "position"

        for _, row in df.iterrows():
            tid = normalize_transcript_id(str(row.get("transcript_id", "")))
            pos = row.get(pos_col)
            if pos is None or pd.isna(pos):
                continue
            pos = int(pos)

            is_cons, state = _classify_row(row, self._fdr_threshold)
            if is_cons:
                self._index[(tid, pos)] = state

    # ---- queries ----------------------------------------------------------

    def lookup(self, transcript_id: str, cds_position_1based: int) -> tuple[bool, str]:
        """Fast lookup: ``(is_conserved, state)``."""
        tid = normalize_transcript_id(transcript_id)
        state = self._index.get((tid, cds_position_1based))
        if state is not None:
            return True, state
        return False, "not_conserved"

    @property
    def conserved_positions(self) -> dict[tuple[str, int], str]:
        """All conserved ``(transcript_id, position)`` -> state."""
        return dict(self._index)

    def __len__(self) -> int:
        return len(self._index)

    def __repr__(self) -> str:
        return f"ConservedPositionIndex({len(self)} conserved positions)"
