# 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

# One-hot lookup: A=0, C=1, G=2, T=3; N/other -> all zeros
BASE_TO_INDEX = {"A": 0, "C": 1, "G": 2, "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.

    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)
    x = np.zeros((1, 4, L), dtype=np.float32)
    for i, ch in enumerate(seq.upper()):
        idx = BASE_TO_INDEX.get(ch)
        if idx is not None:
            x[0, idx, i] = 1.0
    return torch.from_numpy(x).to(device=device, dtype=dtype)
