SilverSight/python/dna_encode_file.py

#!/usr/bin/env python3
"""
dna_encode_file.py — Encode arbitrary files as Hachimoji DNA sequences

Takes any file, chunks it into symbols, builds a monotone LUT,
and outputs the DNA encoding. The LUT makes the encoding self-describing:
anyone with the LUT can decode the DNA back to the original file.

Usage:
    python dna_encode_file.py input.bin [--output output.dna] [--chunk-size 1]
    python dna_encode_file.py --decode input.dna --lut input.lut --output restored.bin

Chunk sizes:
    1 byte  → 256 symbols → 2 bases each (8^2 = 64 ≥ 256? no, need 3 bases)
    1 byte  → 256 symbols → ceil(log8(256)) = 3 bases per symbol
    2 bytes → 65536 symbols → ceil(log8(65536)) = 6 bases per symbol

The monotone LUT assigns shortest DNA sequences to most frequent symbols,
so the DNA output is naturally compact for data with skewed distributions.
"""

from __future__ import annotations

import hashlib
import json
import os
import struct
import sys
import time
from collections import Counter
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass

# ============================================================
# §1  HACHIMOJI ALPHABET (ASCII-ordered for monotone encoding)
# ============================================================

BASES = list("ABCGPSTZ")
N_BASES = len(BASES)
BASE_TO_INDEX = {b: i for i, b in enumerate(BASES)}
INDEX_TO_BASE = {i: b for i, b in enumerate(BASES)}


def int_to_dna(value: int, length: int) -> str:
    """Convert integer to DNA sequence (base-8, fixed length)."""
    seq = []
    for _ in range(length):
        seq.append(INDEX_TO_BASE[value % N_BASES])
        value //= N_BASES
    return "".join(reversed(seq))


def dna_to_int(sequence: str) -> int:
    """Convert DNA sequence to integer."""
    value = 0
    for b in sequence:
        value = value * N_BASES + BASE_TO_INDEX[b]
    return value


def bases_needed(n_symbols: int) -> int:
    """Minimum bases per symbol to represent n_symbols unique keys."""
    if n_symbols <= 0:
        return 0
    length = 1
    while N_BASES ** length < n_symbols:
        length += 1
    return length


# ============================================================
# §2  MONOTONE LUT BUILDER
# ============================================================

def build_frequency_lut(
    data: bytes,
    chunk_size: int = 1,
) -> Tuple[Dict[bytes, str], Dict[str, bytes], int]:
    """Build a monotone LUT from data frequency.

    Most frequent chunks get shortest / lowest-index DNA sequences.
    This is the key to making the encoding work as a compression format:
    frequent symbols → short sequences, rare symbols → long sequences.

    Args:
        data: raw bytes
        chunk_size: bytes per symbol (1, 2, or 4)

    Returns:
        (encode_lut, decode_lut, bases_per_symbol)
    """
    # Count chunk frequencies
    counts: Counter[bytes] = Counter()
    for i in range(0, len(data) - chunk_size + 1, chunk_size):
        chunk = data[i : i + chunk_size]
        if len(chunk) == chunk_size:
            counts[chunk] += 1

    # Handle trailing bytes shorter than chunk_size
    remainder = len(data) % chunk_size
    if remainder > 0:
        counts[data[-remainder:]] += 1

    # Sort by frequency (descending) — most frequent gets lowest index
    sorted_chunks = sorted(counts.keys(), key=lambda c: -counts[c])

    # Build LUT
    n_symbols = len(sorted_chunks)
    bps = bases_needed(n_symbols)

    encode_lut: Dict[bytes, str] = {}
    decode_lut: Dict[str, bytes] = {}

    for rank, chunk in enumerate(sorted_chunks):
        seq = int_to_dna(rank, bps)
        encode_lut[chunk] = seq
        decode_lut[seq] = chunk

    return encode_lut, decode_lut, bps


def build_uniform_lut(chunk_size: int = 1) -> Tuple[Dict[bytes, str], Dict[str, bytes], int]:
    """Build a uniform LUT (all chunks, frequency-independent).

    For chunk_size=1: 256 entries, 3 bases per symbol.
    This is the simple case — no frequency optimization.
    """
    n_symbols = 256 ** chunk_size
    bps = bases_needed(n_symbols)

    encode_lut: Dict[bytes, str] = {}
    decode_lut: Dict[str, bytes] = {}

    for i in range(n_symbols):
        chunk = i.to_bytes(chunk_size, "big")
        seq = int_to_dna(i, bps)
        encode_lut[chunk] = seq
        decode_lut[seq] = chunk

    return encode_lut, decode_lut, bps


# ============================================================
# §3  ENCODING
# ============================================================

@dataclass
class EncodedFile:
    """Result of encoding a file as DNA."""
    dna_sequence: str
    lut: Dict[str, bytes]
    chunk_size: int
    bases_per_symbol: int
    original_size: int
    encoded_length: int
    sha256: str
    encode_time: float

    def to_metadata(self) -> dict:
        return {
            "chunk_size": self.chunk_size,
            "bases_per_symbol": self.bases_per_symbol,
            "original_size": self.original_size,
            "encoded_length": self.encoded_length,
            "sha256": self.sha256,
            "encode_time": round(self.encode_time, 4),
            "lut_entries": len(self.lut),
            "bases": "".join(BASES),
        }



def encode_file(
    data: bytes,
    chunk_size: int = 1,
    use_frequency: bool = True,
) -> EncodedFile:
    """Encode arbitrary bytes as a Hachimoji DNA sequence.

    Args:
        data: raw bytes
        chunk_size: bytes per symbol (1, 2, or 4)
        use_frequency: if True, build frequency-optimized LUT

    Returns:
        EncodedFile with DNA sequence and metadata
    """
    t0 = time.time()

    # Build LUT
    if use_frequency:
        enc_lut, dec_lut, bps = build_frequency_lut(data, chunk_size)
    else:
        enc_lut, dec_lut, bps = build_uniform_lut(chunk_size)

    # Encode
    dna_parts = []
    for i in range(0, len(data), chunk_size):
        chunk = data[i : i + chunk_size]
        if len(chunk) < chunk_size:
            # Pad remainder with zeros
            chunk = chunk + b"\x00" * (chunk_size - len(chunk))
        dna_parts.append(enc_lut[chunk])

    dna = "".join(dna_parts)
    elapsed = time.time() - t0

    return EncodedFile(
        dna_sequence=dna,
        lut=dec_lut,
        chunk_size=chunk_size,
        bases_per_symbol=bps,
        original_size=len(data),
        encoded_length=len(dna),
        sha256=hashlib.sha256(data).hexdigest(),
        encode_time=elapsed,
    )


def decode_file(dna: str, lut: Dict[str, bytes], chunk_size: int, bases_per_symbol: int) -> bytes:
    """Decode a DNA sequence back to bytes using the LUT.

    Args:
        dna: DNA sequence
        lut: decode LUT (sequence → bytes)
        chunk_size: bytes per symbol
        bases_per_symbol: bases per symbol

    Returns:
        Decoded bytes
    """
    parts = []
    for i in range(0, len(dna), bases_per_symbol):
        seq = dna[i : i + bases_per_symbol]
        if seq in lut:
            parts.append(lut[seq])
        else:
            # Unknown sequence — fill with zeros
            parts.append(b"\x00" * chunk_size)

    return b"".join(parts)


# ============================================================
# §4  LUT FILE FORMAT
# ============================================================

def save_lut(lut: Dict[str, bytes], path: str, chunk_size: int, bases_per_symbol: int):
    """Save LUT to a JSON file."""
    data = {
        "format": "hachimoji_lut_v1",
        "bases": "".join(BASES),
        "chunk_size": chunk_size,
        "bases_per_symbol": bases_per_symbol,
        "entries": {seq: chunk.hex() for seq, chunk in lut.items()},
    }
    with open(path, "w") as f:
        json.dump(data, f, indent=2)


def load_lut(path: str) -> Tuple[Dict[str, bytes], int, int]:
    """Load LUT from a JSON file.

    Returns:
        (decode_lut, chunk_size, bases_per_symbol)
    """
    with open(path) as f:
        data = json.load(f)

    chunk_size = data["chunk_size"]
    bps = data["bases_per_symbol"]
    lut = {seq: bytes.fromhex(hex_str) for seq, hex_str in data["entries"].items()}

    return lut, chunk_size, bps


def save_dna(dna: str, path: str):
    """Save DNA sequence to a text file."""
    with open(path, "w") as f:
        f.write(dna)


def load_dna(path: str) -> str:
    """Load DNA sequence from a text file."""
    with open(path) as f:
        return f.read().strip()


# ============================================================
# §5  CLI
# ============================================================

def main():
    import argparse

    parser = argparse.ArgumentParser(description="Encode/decode files as Hachimoji DNA")
    parser.add_argument("input", help="Input file")
    parser.add_argument("--output", "-o", help="Output file (default: input.dna)")
    parser.add_argument("--decode", "-d", action="store_true", help="Decode mode")
    parser.add_argument("--lut", help="LUT file (required for decode)")
    parser.add_argument("--chunk-size", type=int, default=1, choices=[1, 2, 4],
                        help="Bytes per symbol (default: 1)")
    parser.add_argument("--no-frequency", action="store_true",
                        help="Use uniform LUT instead of frequency-optimized")
    parser.add_argument("--save-lut", help="Save LUT to file")

    args = parser.parse_args()

    if args.decode:
        # Decode mode
        if not args.lut:
            parser.error("--lut required for decode mode")
        dna = load_dna(args.input)
        lut, chunk_size, bps = load_lut(args.lut)
        data = decode_file(dna, lut, chunk_size, bps)
        out_path = args.output or args.input.replace(".dna", ".bin")
        with open(out_path, "wb") as f:
            f.write(data)
        print(f"Decoded {len(dna)} bases → {len(data)} bytes → {out_path}")
    else:
        # Encode mode
        with open(args.input, "rb") as f:
            data = f.read()

        result = encode_file(data, args.chunk_size, not args.no_frequency)

        out_path = args.output or args.input + ".dna"
        save_dna(result.dna_sequence, out_path)

        # Save LUT
        lut_path = args.save_lut or args.input + ".lut"
        save_lut(result.lut, lut_path, result.chunk_size, result.bases_per_symbol)

        # Print stats
        m = result.to_metadata()
        print(f"Encoded: {m['original_size']} bytes → {m['encoded_length']} bases")
        print(f"  Chunk size: {m['chunk_size']} byte(s)")
        print(f"  Bases/symbol: {m['bases_per_symbol']}")
        print(f"  LUT entries: {m['lut_entries']}")
        print(f"  SHA-256: {m['sha256']}")
        print(f"  Time: {m['encode_time']}s")
        print(f"  DNA output: {out_path}")
        print(f"  LUT output: {lut_path}")

        # Verify roundtrip
        decoded = decode_file(result.dna_sequence, result.lut, result.chunk_size, result.bases_per_symbol)
        decoded_trimmed = decoded[:len(data)]
        if decoded_trimmed == data:
            print(f"  Roundtrip: ✓ VERIFIED")
        else:
            print(f"  Roundtrip: ✗ FAILED")


if __name__ == "__main__":
    main()