data.py

"""GenBank -> Minerva-ready mixed DNA+AA token sequences. Self-contained (Biopython + numpy)."""

from Bio import SeqIO
from Bio.Seq import Seq
import numpy as np
import re


# ---------------------------------------------------------------------------
# Mixed-token reverse complement
# ---------------------------------------------------------------------------

_DNA_CHARS = frozenset("acgtn")
_COMPLEMENT = str.maketrans("acgtn", "tgcan")
_MARKER_RE = re.compile(r"(<[+-]>)")


def _parse_segments(token_string: str) -> list[tuple[str, str]]:
    """Split a mixed-token string into ``(marker, content)`` pairs.

    >>> _parse_segments("<+>aaa<+>MAQ<->LSH<+>ttt")
    [('<+>', 'aaa'), ('<+>', 'MAQ'), ('<->', 'LSH'), ('<+>', 'ttt')]

    Leading content before any marker is dropped (shouldn't happen in
    well-formed strings).
    """
    parts = _MARKER_RE.split(token_string)
    segments: list[tuple[str, str]] = []
    i = 0
    # Skip any leading text before the first marker
    while i < len(parts) and not _MARKER_RE.fullmatch(parts[i]):
        i += 1
    while i < len(parts):
        marker = parts[i]
        content = parts[i + 1] if i + 1 < len(parts) else ""
        segments.append((marker, content))
        i += 2
    return segments


def _is_dna_segment(content: str) -> bool:
    """Return True if *content* consists entirely of DNA characters (acgtn)."""
    return all(ch in _DNA_CHARS for ch in content)


def rc_mixed_token_string(token_string: str) -> str:
    """Return the reverse-complement of a mixed-modality token string.

    Rules
    -----
    1. Reverse the segment order (3'->5' becomes 5'->3').
    2. Complement DNA segments (a<->t, c<->g, n stays n).
    3. Keep amino-acid segments left-to-right (N->C preserved).
    4. Flip CDS strand markers (``<+>`` <-> ``<->``).
    5. Intergenic (DNA) segments always use ``<+>``.

    Example::

        >>> rc_mixed_token_string("<+>aaa<+>MAQ<->LSH<+>ttt")
        '<+>aaa<+>LSH<->MAQ<+>ttt'
    """
    segments = _parse_segments(token_string)
    if not segments:
        return token_string

    rc_segments: list[str] = []
    for marker, content in reversed(segments):
        if _is_dna_segment(content):
            # DNA: reverse-complement the bases, always use <+>
            rc_segments.append("<+>" + content.translate(_COMPLEMENT)[::-1])
        else:
            # CDS (amino acids): flip the strand marker, keep AA order
            flipped = "<->" if marker == "<+>" else "<+>"
            rc_segments.append(flipped + content)

    return "".join(rc_segments)


def _translate_dna(dna_seq: str, table: int = 11) -> str:
    """Translate DNA sequence to protein, handling codon alignment."""
    remainder = len(dna_seq) % 3
    if remainder != 0:
        dna_seq = dna_seq[:-remainder]

    if len(dna_seq) < 3:
        return ""

    protein_seq = str(Seq(dna_seq).translate(table=table))
    protein_seq = protein_seq.replace("*", "")
    return protein_seq


def _extract_and_translate_cds(
    feature,
    record_seq: str,
    use_existing_translations: bool = False,
    default_translation_table: int = 11,
) -> str:
    """
    Extract and translate a CDS feature with proper handling of:
    - use_existing_translations: Use translation qualifier if present
    - codon_start: Handle reading frame offset
    - Truncated genes: Handle partial 5'/3' genes
    - translation table: honor the feature's ``/transl_table`` qualifier when
      present, otherwise fall back to ``default_translation_table``
      (11 = bacterial/archaeal/plant-plastid, the standard for prokaryotic
      genomes).

    Returns the protein sequence.
    """
    # Use existing translation if available and requested
    if use_existing_translations and "translation" in feature.qualifiers:
        protein_seq = feature.qualifiers["translation"][0]
        return protein_seq.replace("*", "")

    # Extract and translate the CDS sequence
    dna_seq = feature.extract(record_seq)
    if hasattr(dna_seq, 'seq'):
        dna_seq = str(dna_seq.seq)
    else:
        dna_seq = str(dna_seq)

    remainder = len(dna_seq) % 3

    if remainder != 0:
        is_5prime_truncated = False
        is_3prime_truncated = False

        # Check codon_start qualifier
        if "codon_start" in feature.qualifiers:
            codon_start = int(feature.qualifiers["codon_start"][0])
            if codon_start > 1:
                dna_seq = dna_seq[codon_start-1:]
                remainder = len(dna_seq) % 3
                is_5prime_truncated = True

        if remainder != 0:
            location_str = str(feature.location)
            # Check for position operators in string representation
            if '<' in location_str:
                is_5prime_truncated = True
            if '>' in location_str:
                is_3prime_truncated = True

            # Decide how to trim based on truncation status
            if is_3prime_truncated:
                dna_seq = dna_seq[:-remainder]
            elif is_5prime_truncated:
                dna_seq = dna_seq[remainder:]
            else:
                # Default to trimming from the end
                dna_seq = dna_seq[:-remainder]

    if len(dna_seq) < 3:
        return ""

    # Honor a per-CDS /transl_table qualifier when present; otherwise use the
    # caller-supplied default (11 for prokaryotes).
    table = default_translation_table
    if "transl_table" in feature.qualifiers:
        try:
            table = int(feature.qualifiers["transl_table"][0])
        except (ValueError, TypeError):
            table = default_translation_table

    protein_seq = str(Seq(dna_seq).translate(table=table))
    protein_seq = protein_seq.replace("*", "")
    return protein_seq


def _subtract_interval(interval: tuple[int, int], claimed: list[tuple[int, int]]) -> list[tuple[int, int]]:
    """
    Subtract claimed intervals from a single interval.

    Args:
        interval: (start, end) tuple
        claimed: List of (start, end) tuples representing already-claimed regions

    Returns:
        List of (start, end) tuples representing unclaimed portions of the interval
    """
    if not claimed:
        return [interval]

    start, end = interval
    unclaimed = [(start, end)]

    for claim_start, claim_end in claimed:
        new_unclaimed = []
        for seg_start, seg_end in unclaimed:
            # No overlap
            if claim_end <= seg_start or claim_start >= seg_end:
                new_unclaimed.append((seg_start, seg_end))
            else:
                # Overlap - split the segment
                if seg_start < claim_start:
                    new_unclaimed.append((seg_start, claim_start))
                if claim_end < seg_end:
                    new_unclaimed.append((claim_end, seg_end))
        unclaimed = new_unclaimed

    return unclaimed


def _resolve_overlaps(
    features: list[dict],
    record_seq: str,
    mode: str = "expand"
) -> list[dict]:
    """
    Resolve overlapping CDS by prioritizing longer genes.

    Overlaps are resolved regardless of strand (antisense overlaps are treated the same).
    - mode="expand": Keep all CDS features, even overlapping ones
    - mode="long": Shorter overlapping genes are discarded entirely

    Args:
        features: List of feature dicts with 'start', 'end', 'orientation', 'seq', 'gene_name'
        record_seq: The full genomic sequence string
        mode: "expand" to keep all, "long" to discard overlapping genes

    Returns:
        List of features (possibly modified based on mode)
    """
    if not features:
        return features

    # Expand mode: keep all CDS, just sort by position
    if mode == "expand":
        result = [f.copy() for f in features]
        result.sort(key=lambda x: x["start"])
        return result

    # Long mode: keep longest, discard overlapping shorter ones
    # Sort by length (longest first), then by start position
    sorted_by_length = sorted(
        features,
        key=lambda x: (-(x["end"] - x["start"]), x["start"])
    )

    claimed: list[tuple[int, int]] = []
    result = []

    for feat in sorted_by_length:
        # Check if ANY overlap with claimed intervals - if so, discard entirely
        has_overlap = False
        for claim_start, claim_end in claimed:
            if feat["start"] < claim_end and claim_start < feat["end"]:
                has_overlap = True
                break

        if has_overlap:
            continue

        # No overlap - keep the entire feature with its original translation
        result.append(feat.copy())
        claimed.append((feat["start"], feat["end"]))

    # Sort by start position
    result.sort(key=lambda x: x["start"])
    return result


def extract_and_tokenize_gb(
    gb_file: str,
    use_existing_translations: bool = False,
    overlap_mode: str = "expand",
    other_feature_types_to_track: list = None,
    reverse_complement: bool = False,
    translation_table: int = 11,
):
    """
    Extract features from a GenBank file and tokenize them into one sequence per LOCUS.

    Args:
        gb_file: Path to GenBank file
        use_existing_translations: Use translation qualifier if present
        overlap_mode: How to handle overlapping CDS features
            - "expand": Keep all CDS features, even overlapping ones
            - "long": Only keep longest gene at each position, discard overlapping ones
        other_feature_types_to_track: List of additional feature types to extract and track
            (e.g., ['ncRNA', 'tRNA', 'rRNA', 'tmRNA', 'misc_RNA', 'regulatory']).
            Pass ``"all"`` to auto-discover and track every non-CDS, non-source,
            non-gene feature type present in the file.
            These features are tracked but don't affect tokenization. Defaults to None.
        reverse_complement: If True, reverse-complement the mixed-token
            sequence (complement DNA, reverse segment order, flip CDS strand
            markers).  Feature and intergenic lists are reversed accordingly.
        translation_table: Default NCBI genetic-code table used to translate
            CDS features that do not carry their own ``/transl_table``
            qualifier. Defaults to 11 (bacterial/archaeal/plant plastid). A
            feature's own ``/transl_table`` always takes precedence.
    """
    if overlap_mode not in ("expand", "long"):
        raise ValueError(f"Unknown overlap_mode: {overlap_mode}. Must be 'expand' or 'long'")

    # Handle default parameter
    if other_feature_types_to_track is None:
        other_feature_types_to_track = []

    records = list(SeqIO.parse(gb_file, "genbank"))

    # Auto-discover feature types when "all" is requested
    _skip_types = {"CDS", "source", "gene"}
    if other_feature_types_to_track == "all":
        discovered = set()
        for rec in records:
            for feat in rec.features:
                if feat.type not in _skip_types:
                    discovered.add(feat.type)
        other_feature_types_to_track = sorted(discovered)

    tokenized_records = []

    for record_idx, record in enumerate(records):
        sequence = str(record.seq)
        features = []
        other_features = []

        # Extract all CDS features with full translation
        for feature in record.features:
            if feature.type == "CDS":
                start = int(feature.location.start)
                end = int(feature.location.end)
                strand = feature.location.strand
                orientation = True if strand == 1 else False

                # Translate with proper handling of codon_start, truncation, etc.
                protein_seq = _extract_and_translate_cds(
                    feature, record.seq, use_existing_translations,
                    default_translation_table=translation_table,
                )

                if not protein_seq:
                    continue

                # Get gene name if available
                gene_name = ""
                if "gene" in feature.qualifiers:
                    gene_name = feature.qualifiers["gene"][0]
                elif "locus_tag" in feature.qualifiers:
                    gene_name = feature.qualifiers["locus_tag"][0]
                elif "standard_name" in feature.qualifiers:
                    gene_name = feature.qualifiers["standard_name"][0]
                elif "product" in feature.qualifiers:
                    gene_name = feature.qualifiers["product"][0]
                elif "label" in feature.qualifiers:
                    gene_name = feature.qualifiers["label"][0]

                # Get product separately (may differ from gene_name)
                product = ""
                if "product" in feature.qualifiers:
                    product = feature.qualifiers["product"][0]

                features.append({
                    "type": "CDS",
                    "start": start,
                    "end": end,
                    "orientation": orientation,
                    "seq": protein_seq,
                    "gene_name": gene_name,
                    "product": product
                })

            # Extract other feature types if requested
            elif feature.type in other_feature_types_to_track:
                start = int(feature.location.start)
                end = int(feature.location.end)
                strand = feature.location.strand

                # Extract common qualifiers that might be useful
                feature_name = ""
                feature_product = ""
                feature_note = ""

                # Try various qualifier fields for name
                if "gene" in feature.qualifiers:
                    feature_name = feature.qualifiers["gene"][0]
                elif "locus_tag" in feature.qualifiers:
                    feature_name = feature.qualifiers["locus_tag"][0]
                elif "label" in feature.qualifiers:
                    feature_name = feature.qualifiers["label"][0]
                elif "standard_name" in feature.qualifiers:
                    feature_name = feature.qualifiers["standard_name"][0]
                elif "product" in feature.qualifiers:
                    feature_name = feature.qualifiers["product"][0]

                # Try to get product/function information
                if "product" in feature.qualifiers:
                    feature_product = feature.qualifiers["product"][0]
                elif "function" in feature.qualifiers:
                    feature_product = feature.qualifiers["function"][0]

                # Get notes if available
                if "note" in feature.qualifiers:
                    feature_note = feature.qualifiers["note"][0]

                other_features.append({
                    "type": feature.type,
                    "start": start,
                    "end": end,
                    "strand": strand,
                    "name": feature_name,
                    "product": feature_product,
                    "note": feature_note
                })

        # Sort features by position, then by length (longest first for tie-breaking)
        features.sort(key=lambda x: (x["start"], -(x["end"] - x["start"])))

        # Resolve overlapping CDS
        features = _resolve_overlaps(features, sequence, mode=overlap_mode)

        # Process regions in order
        final_tokens = []
        intergenic_regions = []  # Track intergenic regions

        # Handle the case of no CDS features
        if not features:
            if len(sequence) > 0:
                # The entire sequence is intergenic
                final_tokens.append(f"<+>{sequence.lower()}")
                intergenic_regions.append({"start": 0, "end": len(sequence)})
        else:
            # Handle the region before the first CDS
            if features[0]["start"] > 0:
                igs_seq = sequence[0:features[0]["start"]].lower()
                if igs_seq.strip():  # Only add if not empty
                    final_tokens.append(f"<+>{igs_seq}")
                    intergenic_regions.append({"start": 0, "end": features[0]["start"]})

            # Track the "high water mark" - furthest end we've seen
            # This handles overlapping genes correctly
            coverage_end = 0

            # Process each CDS and the intergenic region that follows it
            for i, feature in enumerate(features):
                # Add the CDS with proper orientation
                orientation_token = "<+>" if feature["orientation"] else "<->"
                final_tokens.append(f"{orientation_token}{feature['seq']}")

                # Update high water mark
                coverage_end = max(coverage_end, feature["end"])

                # Add intergenic region after this CDS if not the last one
                if i < len(features) - 1:
                    igs_start = coverage_end
                    igs_end = features[i+1]["start"]
                    if igs_end > igs_start:
                        igs_seq = sequence[igs_start:igs_end].lower()
                        if igs_seq.strip():  # Only add if not empty
                            final_tokens.append(f"<+>{igs_seq}")
                            intergenic_regions.append({"start": igs_start, "end": igs_end})

            # Handle the region after the last CDS
            if coverage_end < len(sequence):
                igs_seq = sequence[coverage_end:].lower()
                if igs_seq.strip():  # Only add if not empty
                    final_tokens.append(f"<+>{igs_seq}")
                    intergenic_regions.append({"start": coverage_end, "end": len(sequence)})

        tokenized_sequence = "".join([token for token in final_tokens if token.strip()])

        # Get the LOCUS name
        locus_name = record.name

        # Map intergenic regions to overlapping other_features
        intergenic_regions_features = []
        for ig_region in intergenic_regions:
            overlapping_features = []
            for other_feat in other_features:
                # Check if feature overlaps with intergenic region
                # Overlap occurs when: feature_start < ig_end AND ig_start < feature_end
                if other_feat["start"] < ig_region["end"] and ig_region["start"] < other_feat["end"]:
                    overlapping_features.append(other_feat)
            intergenic_regions_features.append(overlapping_features)

        # Apply reverse complement if requested
        if reverse_complement:
            tokenized_sequence = rc_mixed_token_string(tokenized_sequence)
            final_tokens = [
                rc_mixed_token_string(tok) for tok in reversed(final_tokens)
            ]
            features = [
                {**f, "orientation": not f["orientation"]}
                for f in reversed(features)
            ]
            intergenic_regions = list(reversed(intergenic_regions))
            intergenic_regions_features = list(reversed(intergenic_regions_features))

        tokenized_records.append({
            "locus_name": locus_name,
            "sequence": tokenized_sequence,
            "tokens": final_tokens,
            "features": features,  # CDS features with gene names
            "other_features": other_features,  # Additional tracked features
            "intergenic_regions": intergenic_regions,  # Intergenic regions
            "intergenic_regions_features": intergenic_regions_features  # Features overlapping each intergenic region
        })

    return tokenized_records