css_modem.py

"""
Acoustic CSS (Chirp Spread Spectrum) modem — desktop reference implementation.
Based on LoRaPHY.m math, adapted for phone speaker/mic frequency range.

Parameters (defaults match a single-chirp MVP):
  fs       = 48000   Hz  — standard AudioRecord/AudioTrack rate
  BW       = 4000    Hz  — chirp sweep width, fits inside 18-22 kHz
  f_center = 20000   Hz  — center of the near-ultrasonic band
  SF       = 7           — spreading factor; 2^SF = 128 samples/symbol, 32 ms
"""

import numpy as np
import scipy.io.wavfile as wav
import scipy.signal as signal
import sys
import os


# ─── Parameters ───────────────────────────────────────────────────────────────

FS = 48_000       # sample rate (Hz)
BW = 4_000        # chirp bandwidth (Hz)
F_CENTER = 20_000 # center frequency (Hz)
SF = 7            # spreading factor — 2^SF samples per symbol

N = 1 << SF       # samples per symbol = 128 at SF7


# ─── Core chirp generation ────────────────────────────────────────────────────

def _base_chirp(up: bool) -> np.ndarray:
    """
    Generate one period of a base up- or down-chirp at baseband.
    Sweeps from -BW/2 to +BW/2 (up) or +BW/2 to -BW/2 (down).
    Returns complex IQ samples, length N.
    """
    t = np.arange(N) / FS
    # instantaneous frequency goes linearly from -BW/2 to +BW/2 over T=N/FS
    # phase = 2π * integral of f(t) = 2π * (-BW/2 * t + BW/(2T) * t^2)
    T = N / FS
    if up:
        phase = 2 * np.pi * (-BW / 2 * t + BW / (2 * T) * t ** 2)
    else:
        phase = 2 * np.pi * (BW / 2 * t - BW / (2 * T) * t ** 2)
    return np.exp(1j * phase)


_UP_CHIRP   = _base_chirp(up=True)
_DOWN_CHIRP = _base_chirp(up=False)


# ─── Modulator ────────────────────────────────────────────────────────────────

def modulate(symbols: list[int]) -> np.ndarray:
    """
    Encode a list of integer symbols (0 … 2^SF-1) into a real audio waveform.
    Prepends 8 up-chirp + 2 down-chirp preamble for sync.
    Returns float32 PCM at FS Hz.
    """
    assert all(0 <= s < N for s in symbols), f"Symbols must be in [0, {N-1}]"

    frames = []

    # Preamble: 8 plain up-chirps + 2 down-chirps (LoRa standard)
    for _ in range(8):
        frames.append(_UP_CHIRP.copy())
    for _ in range(2):
        frames.append(_DOWN_CHIRP.copy())

    # Data symbols: cyclically shift the up-chirp by symbol value
    for s in symbols:
        shifted = np.roll(_UP_CHIRP, s)
        frames.append(shifted)

    # Concatenate all IQ frames, mix up to f_center, take real part
    iq = np.concatenate(frames)
    t = np.arange(len(iq)) / FS
    rf = (iq * np.exp(2j * np.pi * F_CENTER * t)).real

    # Normalize to [-1, 1] and return as float32
    rf /= np.max(np.abs(rf)) + 1e-9
    return rf.astype(np.float32)


# ─── Demodulator ─────────────────────────────────────────────────────────────

def _dechirp_symbol(block: np.ndarray) -> int:
    """
    Demodulate one symbol block (length N complex baseband samples).
    Multiply by conjugate down-chirp (= reference up-chirp*), FFT, find peak bin.
    """
    dechirped = block * np.conj(_UP_CHIRP)
    spectrum  = np.abs(np.fft.fft(dechirped, n=N))
    return int(np.argmax(spectrum))


def _find_preamble(bb: np.ndarray, threshold: float = 0.5) -> int:
    """
    Slide a window of N samples, dechirp, FFT, check if peak is sharp enough
    to be a preamble chirp. Returns the sample offset of the first preamble chirp.
    Returns -1 if no preamble found.
    """
    step = N // 4   # 25% hop — trades CPU for robustness at start

    for start in range(0, len(bb) - N * 10, step):
        block   = bb[start : start + N]
        dechirped = block * np.conj(_UP_CHIRP)
        spectrum  = np.abs(np.fft.fft(dechirped, n=N))
        peak      = np.max(spectrum)
        mean      = np.mean(spectrum)
        if peak / (mean + 1e-9) > threshold * N:
            # Refine: walk forward until the score starts dropping (end of preamble)
            return _refine_sync(bb, start)

    return -1


def _refine_sync(bb: np.ndarray, rough_start: int) -> int:
    """
    Starting from rough_start, walk forward in steps of 1 sample over a small
    window to find the sample-accurate start of the first preamble chirp.
    """
    best_score = -1
    best_pos   = rough_start
    search_range = N // 2

    for offset in range(-search_range, search_range):
        pos = rough_start + offset
        if pos < 0 or pos + N > len(bb):
            continue
        block     = bb[pos : pos + N]
        dechirped = block * np.conj(_UP_CHIRP)
        spectrum  = np.abs(np.fft.fft(dechirped, n=N))
        score     = np.max(spectrum)
        if score > best_score:
            best_score = score
            best_pos   = pos

    return best_pos


def demodulate(pcm: np.ndarray, n_symbols: int) -> list[int]:
    """
    Decode n_symbols data symbols from a float PCM waveform.
    Expects preamble (8 up + 2 down chirps) followed by n_symbols data chirps.
    Returns list of decoded integer symbols.
    """
    # 1. Downconvert to baseband
    t  = np.arange(len(pcm)) / FS
    bb = pcm * np.exp(-2j * np.pi * F_CENTER * t)

    # 2. Low-pass filter to isolate baseband chirp (BW/2 either side)
    nyq    = FS / 2
    cutoff = (BW / 2 + 500) / nyq   # small guard band
    b, a   = signal.butter(4, cutoff, btype='low')
    bb     = signal.lfilter(b, a, bb)

    # 3. Find preamble
    preamble_start = _find_preamble(bb)
    if preamble_start < 0:
        raise RuntimeError("Preamble not found — check signal level and f_center")

    # 4. Skip preamble (8 up + 2 down chirps)
    data_start = preamble_start + N * 10

    # 5. Decode each data symbol
    symbols = []
    for i in range(n_symbols):
        block_start = data_start + i * N
        block_end   = block_start + N
        if block_end > len(bb):
            raise RuntimeError(f"Signal too short for symbol {i}")
        sym = _dechirp_symbol(bb[block_start:block_end])
        symbols.append(sym)

    return symbols


# ─── WAV helpers ─────────────────────────────────────────────────────────────

def save_wav(filename: str, pcm: np.ndarray) -> None:
    """Save float32 PCM as 16-bit WAV at FS."""
    pcm_int16 = (pcm * 32767).astype(np.int16)
    wav.write(filename, FS, pcm_int16)
    duration_ms = len(pcm) / FS * 1000
    print(f"Saved {filename}  ({len(pcm)} samples, {duration_ms:.0f} ms)")


def load_wav(filename: str) -> np.ndarray:
    """Load a WAV file and return float32 PCM. Mono only."""
    rate, data = wav.read(filename)
    assert rate == FS, f"Expected {FS} Hz, got {rate} Hz"
    if data.ndim == 2:
        data = data[:, 0]   # take left channel if stereo
    if data.dtype == np.int16:
        data = data.astype(np.float32) / 32768.0
    return data


# ─── CLI entry point ──────────────────────────────────────────────────────────

def _loopback_test():
    """
    Encode a known symbol list → WAV → decode → verify.
    Run with:  python css_modem.py test
    """
    test_symbols = [0, 1, 63, 64, 127, 42, 7, 100]
    print(f"TX symbols: {test_symbols}")

    pcm = modulate(test_symbols)
    save_wav("test_tx.wav", pcm)

    rx_symbols = demodulate(pcm, len(test_symbols))
    print(f"RX symbols: {rx_symbols}")

    errors = sum(t != r for t, r in zip(test_symbols, rx_symbols))
    print(f"Symbol errors: {errors}/{len(test_symbols)}")
    if errors == 0:
        print("PASS — loopback clean")
    else:
        print("FAIL — check parameters")


def _encode_cmd(args):
    """python css_modem.py encode <sym1> <sym2> ... <out.wav>"""
    if len(args) < 2:
        print("Usage: python css_modem.py encode <sym1> [sym2 ...] <output.wav>")
        sys.exit(1)
    out_file = args[-1]
    symbols  = [int(x) for x in args[:-1]]
    pcm = modulate(symbols)
    save_wav(out_file, pcm)


def _decode_cmd(args):
    """python css_modem.py decode <n_symbols> <in.wav>"""
    if len(args) != 2:
        print("Usage: python css_modem.py decode <n_symbols> <input.wav>")
        sys.exit(1)
    n_symbols = int(args[0])
    in_file   = args[1]
    pcm = load_wav(in_file)
    symbols = demodulate(pcm, n_symbols)
    print(f"Decoded symbols: {symbols}")


if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] == "test":
        _loopback_test()
    elif sys.argv[1] == "encode":
        _encode_cmd(sys.argv[2:])
    elif sys.argv[1] == "decode":
        _decode_cmd(sys.argv[2:])
    else:
        print("Commands: test | encode <syms...> <out.wav> | decode <n> <in.wav>")
        sys.exit(1)