# geney/tis/encoding.py — Sequence encoding for TITER prediction.
from __future__ import annotations

from typing import List

import numpy as np


# ============================================================================
# Sequence Encoding
# ============================================================================

# One-hot encoding scheme for TITER:
# - 8 channels total
# - Channels 0-3: nucleotides OUTSIDE the codon (positions 0-99 and 103-202)
# - Channels 4-7: nucleotides INSIDE the codon (positions 100-102)
# - '$' padding encodes as all zeros

_NUC_TO_IDX = {'A': 0, 'T': 1, 'U': 1, 'C': 2, 'G': 3}


def encode_sequence(seq: str) -> np.ndarray:
    """
    Encode a 203nt sequence for TITER prediction.

    Args:
        seq: 203-character sequence with codon of interest at positions 100-102.
             Use '$' for padding at sequence ends.

    Returns:
        np.ndarray of shape (203, 8)
    """
    if len(seq) != 203:
        raise ValueError(f"Sequence must be exactly 203 nucleotides, got {len(seq)}")

    tensor = np.zeros((203, 8), dtype=np.float32)
    seq = seq.upper()

    for j, nuc in enumerate(seq):
        if nuc == '$':
            continue  # Padding - leave as zeros

        if nuc not in _NUC_TO_IDX:
            continue  # Unknown nucleotide - treat as padding

        idx = _NUC_TO_IDX[nuc]

        # Positions 100-102 are the codon - use channels 4-7
        # All other positions use channels 0-3
        if 100 <= j <= 102:
            tensor[j, idx + 4] = 1.0
        else:
            tensor[j, idx] = 1.0

    return tensor


def encode_sequences(seq_list: List[str]) -> np.ndarray:
    """Encode multiple sequences for batch prediction."""
    return np.stack([encode_sequence(seq) for seq in seq_list])
