# geney/transcription/encoding.py — DNA one-hot encoding for Borzoi.
"""One-hot encode genomic sequences for Borzoi input."""
from __future__ import annotations

import numpy as np
import torch

# Borzoi requires exactly 524,288 bp input
BORZOI_INPUT_LEN = 524_288

# Vectorized lookup table: ord('A')=65 → 0, ord('C')=67 → 1, ord('G')=71 → 2, ord('T')=84 → 3
# All other bytes (including ord('N')=78) map to 255 (sentinel for "no channel").
_ORD_TO_CHANNEL = np.full(256, 255, dtype=np.uint8)
_ORD_TO_CHANNEL[ord("A")] = 0
_ORD_TO_CHANNEL[ord("C")] = 1
_ORD_TO_CHANNEL[ord("G")] = 2
_ORD_TO_CHANNEL[ord("T")] = 3
_ORD_TO_CHANNEL[ord("a")] = 0
_ORD_TO_CHANNEL[ord("c")] = 1
_ORD_TO_CHANNEL[ord("g")] = 2
_ORD_TO_CHANNEL[ord("t")] = 3


def dna_to_onehot(seq: str, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    """Encode a DNA sequence to a ``(1, 4, L)`` one-hot tensor.

    Ambiguous bases (N, etc.) are encoded as all-zero columns.
    Uses vectorized numpy operations for speed (~10x faster than a Python loop
    over 524k characters).

    Args:
        seq: DNA string (should be :data:`BORZOI_INPUT_LEN` bp).
        device: Target torch device.
        dtype: Target torch dtype.

    Returns:
        Tensor of shape ``(1, 4, len(seq))``.
    """
    L = len(seq)
    # Convert string to byte array and look up channel indices
    codes = np.frombuffer(seq.encode("ascii"), dtype=np.uint8)
    channels = _ORD_TO_CHANNEL[codes]  # (L,) with values 0-3 or 255

    # Build one-hot: set x[0, channel, position] = 1 where channel != 255
    x = np.zeros((1, 4, L), dtype=np.float32)
    valid = channels < 4
    positions = np.arange(L, dtype=np.int64)[valid]
    x[0, channels[valid], positions] = 1.0

    return torch.from_numpy(x).to(device=device, dtype=dtype)
