# geney/folding/reference.py — Data-loading utilities for CDS sequences and coordinate maps.
"""Load CDS sequences, CDS-to-genome coordinate maps, and convert positions."""

from __future__ import annotations

import gzip
import pickle
import re
from pathlib import Path

import pandas as pd

__all__ = [
    "normalize_transcript_id",
    "load_cds_sequences",
    "load_mrna_sequences",
    "find_cds_offset",
    "parse_fasta_metadata",
    "select_transcripts",
    "load_cds_genome_map",
    "genome_to_cds_position",
    "cds_to_genome_position",
]


# ---------------------------------------------------------------------------
# Identifier normalisation
# ---------------------------------------------------------------------------

def normalize_transcript_id(tid: str) -> str:
    """Strip Ensembl version suffix (e.g. ``ENST00000123.4`` -> ``ENST00000123``)."""
    return re.sub(r"\.\d+$", "", tid.strip())


# ---------------------------------------------------------------------------
# CDS sequences
# ---------------------------------------------------------------------------

def load_cds_sequences(path: str | Path) -> dict[str, str]:
    """Load CDS sequences keyed by transcript ID.

    Supported formats (auto-detected by extension):

    * **Tabular** (``.txt``, ``.tsv``, ``.csv``, optionally ``.gz``):
      one transcript per line, ``transcript_id<TAB>sequence``.
    * **FASTA** (``.fa``, ``.fasta``, optionally ``.gz``):
      standard ``>header`` / sequence format.

    All transcript IDs are normalised (version stripped).

    Returns:
        ``{transcript_id: cds_dna_sequence}``
    """
    path = Path(path)
    name = path.name.lower()

    if any(name.endswith(ext) for ext in (".fa", ".fasta", ".fa.gz", ".fasta.gz")):
        return _load_fasta(path)
    return _load_tabular_cds(path)


def _open_maybe_gz(path: Path, mode: str = "rt"):
    return gzip.open(path, mode) if str(path).endswith(".gz") else open(path, mode)


def _load_fasta(path: Path) -> dict[str, str]:
    seqs: dict[str, str] = {}
    current_id: str | None = None
    chunks: list[str] = []
    with _open_maybe_gz(path) as fh:
        for line in fh:
            line = line.strip()
            if line.startswith(">"):
                if current_id is not None:
                    seqs[normalize_transcript_id(current_id)] = "".join(chunks).upper()
                current_id = line[1:].split()[0]
                chunks = []
            else:
                chunks.append(line)
    if current_id is not None:
        seqs[normalize_transcript_id(current_id)] = "".join(chunks).upper()
    return seqs


def _load_tabular_cds(path: Path) -> dict[str, str]:
    seqs: dict[str, str] = {}
    with _open_maybe_gz(path) as fh:
        for line in fh:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split("\t") if "\t" in line else line.split()
            if len(parts) >= 2:
                seqs[normalize_transcript_id(parts[0])] = parts[1].upper().strip()
    return seqs


# ---------------------------------------------------------------------------
# Mature mRNA (cDNA) sequences
# ---------------------------------------------------------------------------

def load_mrna_sequences(path: str | Path) -> dict[str, str]:
    """Load mature mRNA (cDNA) sequences keyed by transcript ID.

    This is a thin wrapper around the same FASTA / tabular loader used by
    :func:`load_cds_sequences`, but named explicitly for the mature-mRNA
    use-case (5'UTR + CDS + 3'UTR, typically from ``cdna.all.fa.gz``).

    Args:
        path: Path to a FASTA or tabular file (optionally gzipped).

    Returns:
        ``{transcript_id: mrna_dna_sequence}``
    """
    path = Path(path)
    name = path.name.lower()

    if any(name.endswith(ext) for ext in (".fa", ".fasta", ".fa.gz", ".fasta.gz")):
        return _load_fasta(path)
    return _load_tabular_cds(path)


def find_cds_offset(mrna_seq: str, cds_seq: str) -> int:
    """Locate the CDS as a substring of the mature mRNA.

    Returns the 0-based start index of *cds_seq* within *mrna_seq*,
    which equals the 5'UTR length.

    Some Ensembl CDS entries have leading/trailing ``N`` bases where the
    exact sequence is unknown.  If the raw CDS is not found, the function
    retries after stripping those ``N`` characters.

    Raises:
        ValueError: If *cds_seq* is not found within *mrna_seq* even after
            stripping ambiguous bases.
    """
    mrna_upper = mrna_seq.upper()
    cds_upper = cds_seq.upper()

    idx = mrna_upper.find(cds_upper)
    if idx != -1:
        return idx

    # Retry with leading/trailing N's stripped
    stripped = cds_upper.strip("N")
    if stripped and stripped != cds_upper:
        idx = mrna_upper.find(stripped)
        if idx != -1:
            return idx

    raise ValueError(
        f"CDS not found as substring of mRNA "
        f"(CDS len={len(cds_seq)}, mRNA len={len(mrna_seq)})"
    )


# ---------------------------------------------------------------------------
# FASTA metadata and transcript selection
# ---------------------------------------------------------------------------

def parse_fasta_metadata(path: str | Path) -> dict[str, dict]:
    """Parse Ensembl FASTA headers into per-transcript metadata.

    Extracts ``gene_id``, ``gene_biotype``, ``transcript_biotype``,
    ``gene_symbol``, ``seq_length``, and ``chromosome`` from standard
    Ensembl header lines.

    Args:
        path: Path to an Ensembl FASTA file (CDS or cDNA, optionally gzipped).

    Returns:
        ``{transcript_id: {gene_id, gene_biotype, transcript_biotype,
        gene_symbol, seq_length, chromosome}}``
    """
    path = Path(path)
    meta: dict[str, dict] = {}
    current_id: str | None = None
    current_header: str = ""
    seq_len = 0

    def _flush():
        if current_id is None:
            return
        info: dict = {"seq_length": seq_len}
        for tag in ("gene", "gene_biotype", "transcript_biotype", "gene_symbol"):
            m = re.search(rf"{tag}:(\S+)", current_header)
            if m:
                val = m.group(1)
                key = "gene_id" if tag == "gene" else tag
                info[key] = normalize_transcript_id(val) if tag == "gene" else val
        # chromosome from the coordinate field
        coord = re.search(r"(?:chromosome|scaffold):(\S+)", current_header)
        if coord:
            parts = coord.group(1).split(":")
            info["chromosome"] = parts[1] if len(parts) > 1 else parts[0]
        meta[normalize_transcript_id(current_id)] = info

    with _open_maybe_gz(path) as fh:
        for line in fh:
            line = line.strip()
            if line.startswith(">"):
                _flush()
                current_id = line[1:].split()[0]
                current_header = line
                seq_len = 0
            else:
                seq_len += len(line)
    _flush()
    return meta


def select_transcripts(
    cds_metadata: dict[str, dict],
    level: str = "one_per_gene",
    principal_transcripts: set[str] | list[str] | None = None,
    top_n: int = 3,
) -> list[str]:
    """Select a subset of transcripts for precomputation.

    Args:
        cds_metadata: Output of :func:`parse_fasta_metadata` on the CDS FASTA.
        level: Selection strategy:

            - ``"one_per_gene"`` — one transcript per gene. Prefers a
              principal transcript (if *principal_transcripts* provided),
              otherwise picks the longest CDS. **~23K transcripts.**
            - ``"top_n_per_gene"`` — top *top_n* by CDS length per gene.
              **~60K at top_n=3.**
            - ``"all_protein_coding"`` — every protein-coding transcript.
              **~222K transcripts.**

        principal_transcripts: Optional set of transcript IDs considered
            "principal" (e.g. MANE Select).  Used by ``"one_per_gene"``
            to prefer known principal isoforms over longest.
        top_n: Number of transcripts per gene for ``"top_n_per_gene"``.

    Returns:
        Sorted list of transcript IDs.
    """
    if principal_transcripts is not None:
        principal_transcripts = {
            normalize_transcript_id(t) for t in principal_transcripts
        }

    # Group protein-coding transcripts by gene
    gene_transcripts: dict[str, list[tuple[str, int]]] = {}
    for tid, info in cds_metadata.items():
        if info.get("transcript_biotype") != "protein_coding":
            continue
        gid = info.get("gene_id", tid)
        gene_transcripts.setdefault(gid, []).append(
            (tid, info.get("seq_length", 0))
        )

    # Sort each gene's transcripts: principal first, then by length descending
    for gid in gene_transcripts:
        gene_transcripts[gid].sort(
            key=lambda x: (
                0 if principal_transcripts and x[0] in principal_transcripts else 1,
                -x[1],
            )
        )

    selected: list[str] = []
    if level == "one_per_gene":
        for gid, tlist in gene_transcripts.items():
            selected.append(tlist[0][0])
    elif level == "top_n_per_gene":
        for gid, tlist in gene_transcripts.items():
            selected.extend(tid for tid, _ in tlist[:top_n])
    elif level == "all_protein_coding":
        for gid, tlist in gene_transcripts.items():
            selected.extend(tid for tid, _ in tlist)
    else:
        raise ValueError(
            f"Unknown level {level!r}. "
            f"Choose from: one_per_gene, top_n_per_gene, all_protein_coding"
        )

    return sorted(selected)


# ---------------------------------------------------------------------------
# CDS ↔ genome coordinate map
# ---------------------------------------------------------------------------

def load_cds_genome_map(path: str | Path) -> dict[str, list[int]]:
    """Load CDS-position ↔ genome-position mapping.

    Supported formats:

    * **Pickle** (``.pickle``, ``.pkl``):
      ``dict[transcript_id -> <mapping>]`` where mapping is one of:

      - ``list[int]``: genome positions indexed by CDS position (0-based).
      - ``dict[int, int]``: ``{cds_pos: genome_pos}``.
      - ``dict`` with key ``'positions'``: e.g. ``{'chrom': ..., 'positions': [...]}``.

    * **Parquet** (``.parquet``):
      columns ``transcript_id``, ``cds_pos``, ``genome_pos``.

    Returns:
        ``{transcript_id: [genome_pos_for_cds_0, genome_pos_for_cds_1, ...]}``
    """
    path = Path(path)
    name = path.name.lower()

    if name.endswith(".parquet"):
        return _load_map_parquet(path)
    return _load_map_pickle(path)


def _load_map_pickle(path: Path) -> dict[str, list[int]]:
    with open(path, "rb") as fh:
        raw = pickle.load(fh)

    result: dict[str, list[int]] = {}
    for tid, value in raw.items():
        norm = normalize_transcript_id(str(tid))
        if isinstance(value, dict):
            if "positions" in value:
                # {chrom: ..., positions: [...], ...}
                result[norm] = list(value["positions"])
            else:
                # {cds_pos: genome_pos, ...}
                if value:
                    max_key = max(value.keys())
                    result[norm] = [value.get(i) for i in range(max_key + 1)]
                else:
                    result[norm] = []
        elif isinstance(value, (list, tuple)):
            result[norm] = list(value)
        else:
            # Single value or unknown — store as-is wrapped in list
            result[norm] = [value] if value is not None else []
    return result


def _load_map_parquet(path: Path) -> dict[str, list[int]]:
    df = pd.read_parquet(path)
    result: dict[str, list[int]] = {}
    for tid, grp in df.groupby("transcript_id"):
        sorted_grp = grp.sort_values("cds_pos")
        result[normalize_transcript_id(str(tid))] = sorted_grp["genome_pos"].tolist()
    return result


# ---------------------------------------------------------------------------
# Position conversion
# ---------------------------------------------------------------------------

_COMPLEMENT = str.maketrans("ACGTacgt", "TGCAtgca")


def complement(base: str) -> str:
    """Return the Watson–Crick complement of a single base."""
    return base.translate(_COMPLEMENT)


def genome_to_cds_position(
    cds_map: dict[str, list[int]],
    transcript_id: str,
    genome_pos: int,
) -> int | None:
    """Convert a genome position to 0-based CDS position.

    Returns ``None`` if the transcript or position is not found.
    """
    positions = cds_map.get(normalize_transcript_id(transcript_id))
    if positions is None:
        return None
    try:
        return positions.index(genome_pos)
    except ValueError:
        return None


def cds_to_genome_position(
    cds_map: dict[str, list[int]],
    transcript_id: str,
    cds_pos_0based: int,
) -> int | None:
    """Convert a 0-based CDS position back to a genome position.

    Returns ``None`` if the transcript or position is not found.
    """
    positions = cds_map.get(normalize_transcript_id(transcript_id))
    if positions is None or cds_pos_0based < 0 or cds_pos_0based >= len(positions):
        return None
    return positions[cds_pos_0based]
