bandsaunter/bandsaunter/morse.py

"""CW / Morse detection and decoding.

Takes demodulated audio (or a complex baseband block), finds the keyed tone,
recovers the element timing without being told the speed, and decodes the
text.  Also reports a confidence so the scanner can tell real Morse from any
other on-off keyed carrier.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field

import numpy as np
from scipy import signal as sps

__all__ = ["decode_morse", "find_morse", "MorseResult", "MORSE_TABLE",
           "encode_morse", "SHORT_TONE_DB", "MIN_SAMPLES"]


MORSE_TABLE: dict[str, str] = {
    ".-": "A", "-...": "B", "-.-.": "C", "-..": "D", ".": "E",
    "..-.": "F", "--.": "G", "....": "H", "..": "I", ".---": "J",
    "-.-": "K", ".-..": "L", "--": "M", "-.": "N", "---": "O",
    ".--.": "P", "--.-": "Q", ".-.": "R", "...": "S", "-": "T",
    "..-": "U", "...-": "V", ".--": "W", "-..-": "X", "-.--": "Y",
    "--..": "Z",
    "-----": "0", ".----": "1", "..---": "2", "...--": "3", "....-": "4",
    ".....": "5", "-....": "6", "--...": "7", "---..": "8", "----.": "9",
    ".-.-.-": ".", "--..--": ",", "..--..": "?", ".----.": "'",
    "-.-.--": "!", "-..-.": "/", "-.--.": "(", "-.--.-": ")",
    ".-...": "&", "---...": ":", "-.-.-.": ";", "-...-": "=",
    ".-.-.": "+", "-....-": "-", "..--.-": "_", ".-..-.": '"',
    "...-..-": "$", ".--.-.": "@",
    # Common prosigns, rendered between angle brackets.
    "...-.-": "<SK>", "-.-.-": "<KA>", ".-.-": "<AA>", "...-.": "<SN>",
    "-...-.-": "<BK>", "........": "<HH>", ".-.-.-.-": "<ERR>",
}

_REVERSE = {v: k for k, v in MORSE_TABLE.items() if len(v) == 1}

# Morse timing is defined against PARIS = 50 dot units per word.
_DOTS_PER_WORD = 50.0

# How far the keyed tone has to stand above the rest of its band before a
# decode of only a character or two is believed.  See MorseResult.is_morse.
SHORT_TONE_DB = 20.0

# What _find_tone needs before it can measure anything: one FFT's worth of
# samples.  It is also the floor on the whole decode, because a station that
# identifies itself in Morse and nothing else is on the air for well under a
# second, and a floor set for comfort throws those away unheard.
MIN_SAMPLES = 1024

# Silence of this many dot units or more separates words.  ITU says seven;
# five is where the decoder splits them, because a hand key stretches the
# short gaps and shortens the long ones.
_WORD_GAP_UNITS = 5.0

# A mark this many dots long is not a Morse element.  A dash is three, and
# generous slop puts the longest believable one at four or five; beyond that
# the key is not down for a reason the code is reading, and it is almost
# always the transmission the ident was sent over -- a repeater's tail, a
# steady tone, the voice that came before.
_MAX_ELEMENT_UNITS = 5.0


@dataclass
class MorseResult:
    text: str = ""
    wpm: float = 0.0
    tone_hz: float = 0.0
    dot_seconds: float = 0.0
    confidence: float = 0.0
    n_elements: int = 0
    n_characters: int = 0
    timing_fit: float = 0.0
    undecoded: int = 0
    snr_db: float = 0.0
    head_cut: bool = False        # the capture opened mid-character
    tail_cut: bool = False        # and/or closed mid-character
    notes: list[str] = field(default_factory=list)

    @property
    def complete_text(self) -> str:
        """The part of the text that is certainly what was sent.

        The character sliced by the capture window has already been dropped,
        but what is left of the word it was in can still read as a whole one:
        "K1AA" cut short is "K1A", which is somebody else's callsign, and
        this text is looked at for callsigns.  So a word touching a cut end
        is not reported either.  What comes back is what a station can be
        identified from; ``text`` remains everything that was read.
        """
        if not (self.head_cut or self.tail_cut):
            return self.text
        words = self.text.split()
        if self.head_cut and words:
            words = words[1:]
        if self.tail_cut and words:
            words = words[:-1]
        return " ".join(words)

    @property
    def is_morse(self) -> bool:
        """Whether this decode should be believed.

        Two bars, not one.  A few seconds of text clears the ordinary one.
        Anything shorter has to be cleaner than that, because at two or three
        characters there is not enough of it for a wrong reading to
        contradict itself -- but shorter is exactly what a station giving
        nothing but its callsign sends, and a bar it could never clear would
        throw away the transmissions most worth having.

        Below three elements there is nothing left to be right about: one
        keyed pulse is an E or a T whether a person sent it or the squelch
        opened on a click, so that is where this stops.
        """
        # Nobody sends below 5 or above 60 WPM.  A "decode" outside that
        # range is the timing estimator latching onto something that is not
        # Morse -- speech syllables, for instance.
        if not 5.0 <= self.wpm <= 60.0:
            return False
        if self.undecoded > 0.3 * max(1, self.n_characters):
            return False
        # Every character one element long means every character is an E or a
        # T, which is what a run of unstructured pulses always decodes to: it
        # can hardly be wrong, because there is nothing in it to get wrong.
        # A capture of noise on 445.5 MHz came back as "T T T E E E E E E E
        # E" with the element mix and the timing fit both inside their bands,
        # and no station has ever identified itself that way.
        if self.n_characters and self.n_elements <= self.n_characters:
            return False
        if self.n_characters >= 3 and self.n_elements >= 8:
            return self.confidence >= 0.5 and self.timing_fit >= 0.7
        # Timing alone is not enough down here.  With three or four elements
        # the dot length is fitted to those very elements, so they land on the
        # grid whatever produced them -- a third of a second of white noise
        # decodes as a perfectly timed V.  So a short burst has to have been
        # a tone as well: measured over two hundred noise blocks the loudest
        # bin never rose 13 dB above the median of its band, while keying at
        # 3 dB SNR sits above 40.  Twenty is between the two with room to
        # spare on both sides.
        return (self.n_characters >= 1 and self.n_elements >= 3
                and self.timing_fit >= 0.95 and self.undecoded == 0
                and self.snr_db >= SHORT_TONE_DB
                and self.confidence >= 0.4)

    def summary(self) -> str:
        if not self.text:
            return "keyed carrier, no readable Morse"
        return f"{self.wpm:.0f} WPM: {self.text.strip()}"


def encode_morse(text: str) -> str:
    """Encode text to Morse -- handy for generating test signals."""
    out = []
    for ch in text.upper():
        if ch == " ":
            out.append("/")
        elif ch in _REVERSE:
            out.append(_REVERSE[ch])
    return " ".join(out)


# ---------------------------------------------------------------------------

def _tone_candidates(audio: np.ndarray, fs: float, most: int = 3,
                     lo: float = 200.0,
                     hi: float = 3000.0) -> list[tuple[float, float]]:
    """The narrow tones in a clip, strongest first, as ``(hz, prominence_db)``.

    Averaged across the whole clip rather than measured on the front of it.
    A station that idents in Morse does it once, wherever in the capture it
    happens to fall, and one transform over the opening seconds cannot see a
    tone that had not started yet -- on a half-minute capture with the ident
    two-thirds of the way through, it read the harmonic of something else
    and the ident was never decoded at all.

    More than one candidate because the loudest tone is not always the keyed
    one: a CTCSS tone, a data subcarrier or a carrier's own whine can all be
    steadier and stronger than the ident sent over the top of them.
    """
    n = 1 << int(math.floor(math.log2(max(1024, min(audio.size, 1 << 14)))))
    if audio.size < n:
        return []
    x = np.asarray(audio, dtype=np.float64)
    win = np.hanning(n)
    step = max(1, n // 2)
    acc = np.zeros(n // 2 + 1)
    frames = 0
    for at in range(0, x.size - n + 1, step):
        block = x[at:at + n]
        acc += np.abs(np.fft.rfft((block - block.mean()) * win, n))
        frames += 1
    if not frames:
        return []
    spec = acc / frames
    freqs = np.fft.rfftfreq(n, 1.0 / fs)
    band = (freqs >= lo) & (freqs <= min(hi, fs / 2.2))
    if not np.any(band):
        return []
    sub, subf = spec[band], freqs[band]
    floor = float(np.median(sub)) + 1e-12
    # One candidate per peak: neighbouring bins of the same tone are the same
    # tone, and the band-pass that follows is 300 Hz wide anyway.
    out: list[tuple[float, float]] = []
    for k in np.argsort(sub)[::-1]:
        if any(abs(subf[k] - hz) < 150.0 for hz, _ in out):
            continue
        out.append((float(subf[k]),
                    20.0 * math.log10((sub[k] + 1e-12) / floor)))
        if len(out) >= most:
            break
    return out


def _find_tone(audio: np.ndarray, fs: float,
               lo: float = 200.0, hi: float = 3000.0) -> tuple[float, float]:
    """Locate the keyed tone.  Returns ``(freq_hz, prominence_db)``."""
    found = _tone_candidates(audio, fs, most=1, lo=lo, hi=hi)
    return found[0] if found else (0.0, 0.0)


def _tone_envelope(audio: np.ndarray, fs: float, tone_hz: float,
                   width: float = 150.0) -> tuple[np.ndarray, float]:
    """Band-pass around the tone and return its envelope at a reduced rate."""
    nyq = fs / 2.0
    lo = max(20.0, tone_hz - width) / nyq
    hi = min(0.99, (tone_hz + width) / nyq)
    if hi <= lo:
        return np.abs(audio), fs
    taps = sps.firwin(255, [lo, hi], pass_zero=False)
    filtered = sps.lfilter(taps, [1.0], audio)[255:]
    env = np.abs(sps.hilbert(filtered))

    # 500 Hz of envelope bandwidth resolves elements down to ~2 ms.
    dec = max(1, int(fs // 1000))
    if dec > 1:
        env = sps.decimate(env, dec, ftype="fir", zero_phase=True)
    return env, fs / dec


def _threshold(env: np.ndarray) -> float:
    """Split the envelope into key-down and key-up populations."""
    lo = float(np.percentile(env, 15))
    hi = float(np.percentile(env, 85))
    if hi <= lo:
        return float(np.median(env))
    # Midway in the log domain: keying contrast is multiplicative, not additive.
    return float(math.sqrt(max(lo, 1e-12) * max(hi, 1e-12)))


def _despeckle(on: np.ndarray, min_len: int) -> np.ndarray:
    """Remove key transitions too short to be real elements.

    Filter ring at the edges of each element crosses the threshold for a
    sample or two.  Left in, those specks become the shortest "element" and
    drag the dot-length estimate down to nothing, so flatten anything below
    the fastest plausible keying speed into its neighbours.
    """
    if min_len <= 1 or on.size == 0:
        return on
    out = on.copy()
    for _ in range(4):
        changed = False
        for state, start, length in [(s, a, b) for s, a, b in _run_spans(out)]:
            if length < min_len:
                out[start:start + length] = not state
                changed = True
        if not changed:
            break
    return out


def _run_spans(on: np.ndarray):
    if on.size == 0:
        return
    change = np.flatnonzero(np.diff(on.astype(np.int8)))
    bounds = np.concatenate(([0], change + 1, [on.size]))
    for i in range(bounds.size - 1):
        yield bool(on[bounds[i]]), int(bounds[i]), int(bounds[i + 1] - bounds[i])


def _runs(on: np.ndarray) -> list[tuple[bool, int]]:
    if on.size == 0:
        return []
    change = np.flatnonzero(np.diff(on.astype(np.int8)))
    bounds = np.concatenate(([0], change + 1, [on.size]))
    return [(bool(on[bounds[i]]), int(bounds[i + 1] - bounds[i]))
            for i in range(bounds.size - 1)]


def _estimate_dot(on_lengths: list[int], off_lengths: list[int]) -> float:
    """Two-means clustering in the log domain gives a speed-agnostic dot length."""
    vals = np.array([v for v in on_lengths if v > 0], dtype=np.float64)
    if vals.size == 0:
        return 0.0
    if vals.size < 4:
        return float(np.min(vals))

    logs = np.log(vals)
    lo, hi = float(logs.min()), float(logs.max())
    if hi - lo < 0.35:          # all one length -- everything is a dot
        return float(np.median(vals))
    c0, c1 = lo, hi
    for _ in range(30):
        d0 = np.abs(logs - c0)
        d1 = np.abs(logs - c1)
        g0, g1 = logs[d0 <= d1], logs[d1 < d0]
        if g0.size == 0 or g1.size == 0:
            break
        n0, n1 = float(g0.mean()), float(g1.mean())
        if abs(n0 - c0) < 1e-6 and abs(n1 - c1) < 1e-6:
            c0, c1 = n0, n1
            break
        c0, c1 = n0, n1
    dot = math.exp(c0)
    dash = math.exp(c1)
    # A dash is three dots; if the two clusters are close to that ratio, use
    # both to refine the estimate.
    if 2.0 <= dash / max(dot, 1e-9) <= 4.5:
        dot = 0.5 * (dot + dash / 3.0)
    # The inter-element gap is also one dot long -- a useful cross-check.
    offs = np.array([v for v in off_lengths if v > 0], dtype=np.float64)
    if offs.size >= 3:
        short = offs[offs <= np.percentile(offs, 40)]
        if short.size and 0.5 * dot <= np.median(short) <= 2.0 * dot:
            dot = 0.5 * (dot + float(np.median(short)))
    return dot


# How long a stretch of audio to hand the decoder at a time when hunting for
# an ident inside a longer capture, and how far to slide between tries.  An
# ident is a callsign at 15-25 WPM -- two to six seconds -- and the lengths
# below bracket that with room either side.  The step is a third of the
# window rather than half because the window has to fall *around* the ident,
# not merely overlap it: everything the decoder measures, the tone and the
# key-down threshold both, is measured over the whole window, so a window
# that is mostly something else measures that instead.  Stepping by half a
# window found neither of the two idents in a night of recordings; stepping
# by a third found both.
FIND_WINDOWS = (4.0, 7.0, 12.0)
FIND_STEP_FRACTION = 1.0 / 3.0
FIND_MIN_STEP = 1.0


def find_morse(audio: np.ndarray, sample_rate: float) -> MorseResult | None:
    """Hunt for a Morse ident anywhere in a capture.

    :func:`decode_morse` reads a clip that *is* Morse.  This looks for one
    inside a clip that is mostly something else -- the case that matters on
    the land-mobile bands, where a base station idents in CW over the top of
    its own carrier and the rest of the capture is voice, noise, or a steady
    tone.  Handed the whole of such a capture the decoder has no chance: it
    picks its tone and its key-down threshold from the whole window, and on a
    half-minute recording with five seconds of keying in the middle both come
    out of the other twenty-five.

    So the same decoder is offered a series of shorter windows and the
    reading that identifies a station best is kept.  Nothing is loosened to
    make that work: every window is judged by :attr:`MorseResult.is_morse`
    exactly as a whole capture would be.  Across 677 real captures from two
    nights of scanning it claimed Morse in two, and both were idents.

    Returns None when no window read as Morse.
    """
    audio = np.asarray(audio, dtype=np.float64).ravel()
    if audio.size < MIN_SAMPLES:
        return None
    rate = float(sample_rate)
    spans: set[tuple[int, int]] = set()
    for seconds in FIND_WINDOWS:
        width = int(seconds * rate)
        if width < MIN_SAMPLES:
            continue
        if audio.size <= width:
            spans.add((0, audio.size))
            continue
        step = max(1, int(max(FIND_MIN_STEP,
                              seconds * FIND_STEP_FRACTION) * rate))
        for at in range(0, audio.size - width + 1, step):
            spans.add((at, at + width))
        spans.add((audio.size - width, audio.size))
    # The whole capture, always.  When the capture *is* Morse -- a beacon
    # keying continuously through it -- the longest window is the best one,
    # because a word is only certain when a gap bounds it at both ends and a
    # short window may not contain one.  Leaving this out lost a beacon the
    # decoder had always read.
    spans.add((0, audio.size))

    best: MorseResult | None = None
    for lo, hi in sorted(spans):
        try:
            found = decode_morse(audio[lo:hi], rate)
        except Exception:
            continue
        if not found.is_morse:
            continue
        # The longest identifiable reading wins.  complete_text rather than
        # text, because a window that clipped the ident reports fewer
        # characters it can stand behind, which is exactly the ranking
        # wanted: the window that fell around the ident beats the ones that
        # fell across it.
        if best is None or len(found.complete_text) > len(best.complete_text) \
                or (len(found.complete_text) == len(best.complete_text)
                    and found.confidence > best.confidence):
            best = found
    return best


def decode_morse(audio: np.ndarray, sample_rate: float,
                 min_elements: int = 3) -> MorseResult:
    """Decode CW from a block of demodulated audio.

    ``audio`` should be real audio containing the beat note (what the ``cw``
    demodulator produces).  Speed is estimated from the signal itself, so no
    WPM setting is needed.

    ``min_elements`` counts key transitions, up and down together, after the
    partial runs at each end have been dropped.  Three is one character with
    structure of its own -- a K, an R, a digit -- which is the shortest thing
    that can be told from a click.
    """
    res = MorseResult()
    audio = np.asarray(audio, dtype=np.float64).ravel()
    if audio.size < MIN_SAMPLES:
        res.notes.append("too short to decode")
        return res

    tone, prom = _find_tone(audio, sample_rate)
    res.tone_hz = tone
    res.snr_db = prom
    if tone <= 0 or prom < 6.0:
        res.notes.append("no steady tone found")
        return res

    env, fs_env = _tone_envelope(audio, sample_rate, tone)
    if env.size < 32:
        res.notes.append("envelope too short")
        return res

    thr = _threshold(env)
    on = env > thr
    # 6 ms is a dot at about 200 WPM -- far beyond any real operator or
    # machine, so anything shorter is filter ring, not keying.
    on = _despeckle(on, max(2, int(0.006 * fs_env)))
    runs = _runs(on)
    # Drop the leading and trailing partial runs -- they are cut off by the
    # capture window and would corrupt the timing estimate.  Keep them,
    # though: once the dot length is known they say whether the character at
    # each end is cut off too, and half a character decodes to a different
    # one rather than to less of the same.  A K missing its first dash is an
    # A; a W missing its first dot is an M.
    edges: tuple = (None, None)
    if len(runs) >= 3:
        edges = (runs[0], runs[-1])
        runs = runs[1:-1]
    if len(runs) < min_elements:
        res.notes.append("not enough keying transitions")
        return res

    on_lengths = [n for state, n in runs if state]
    off_lengths = [n for state, n in runs if not state]
    if not on_lengths:
        res.notes.append("carrier never keys down")
        return res

    dot = _estimate_dot(on_lengths, off_lengths)
    if dot <= 0:
        res.notes.append("could not estimate element length")
        return res
    res.dot_seconds = dot / fs_env
    if not (0.008 <= res.dot_seconds <= 0.5):
        res.notes.append(f"implausible element length ({res.dot_seconds*1e3:.0f} ms)")
        return res
    res.wpm = 1.2 / res.dot_seconds * (_DOTS_PER_WORD / 50.0)

    def _cut(edge) -> bool:
        """Whether the character at this end of the window is incomplete.

        Three ways for it to go.  Key-down at the boundary is the obvious
        cut: the element itself is sliced.  Silence too short to be a word
        gap is the other, and the one that is easy to miss -- the window
        opened partway through a word, so the rest of that word is outside
        it, and what is left of it can read as a whole word of its own.

        The third is key-down for far longer than any element lasts, and it
        is not a cut at all.  A station that idents in Morse over an FM
        carrier leaves a steady tone either side of the ident; the window
        opens in the middle of that tone, and nothing was sliced, because
        nothing was being keyed.  Reading it as a truncated character threw
        away the first and last letter of every such ident -- and with them,
        by way of ``complete_text``, the whole callsign, which is one word
        with no gap in it to survive the drop.
        """
        if edge is None:
            return False
        state, length = edge
        if state:
            return (length / dot) <= _MAX_ELEMENT_UNITS
        return (length / dot) < _WORD_GAP_UNITS

    head_cut, tail_cut = _cut(edges[0]), _cut(edges[1])
    res.head_cut, res.tail_cut = head_cut, tail_cut

    # ---- element decision --------------------------------------------
    # Characters are carried with the number of elements each one took, so
    # that dropping a character cut off by the capture window drops its
    # elements with it and the counts stay true.
    decoded: list[tuple[str, int]] = []
    current = ""

    def flush():
        nonlocal current
        if not current:
            return
        decoded.append((MORSE_TABLE.get(current, "<?>"), len(current)))
        current = ""

    for state, n in runs:
        units = n / dot
        if state:
            current += "." if units < 2.0 else "-"
        else:
            if units < 2.0:
                continue                    # gap between elements of a letter
            flush()
            if units >= _WORD_GAP_UNITS:
                decoded.append((" ", 0))    # word gap
    flush()

    # The characters at a truncated end are incomplete, so they are dropped
    # rather than reported: a partial character is not a smaller reading of
    # what was sent, it is a different one, and this text is looked at for
    # callsigns.  A station heard through half its ident is better reported
    # as half an ident than as a different station.
    if head_cut and decoded:
        decoded = decoded[1:]
        res.notes.append("first character was cut off and dropped")
    if tail_cut and decoded:
        decoded = decoded[:-1]
        res.notes.append("last character was cut off and dropped")
    text_parts = [ch for ch, _ in decoded]

    text = "".join(text_parts)
    # Collapse runs of spaces the timing may have produced.
    text = " ".join(text.split(" ")).strip() if text else ""
    undecoded = sum(1 for ch, _ in decoded if ch == "<?>")
    res.text = text
    res.n_elements = sum(n for _, n in decoded)
    res.n_characters = sum(1 for ch, _ in decoded if ch != " ")
    res.undecoded = undecoded

    # ---- confidence ----------------------------------------------------
    if res.n_characters == 0:
        res.confidence = 0.0
        res.notes.append("no characters resolved")
        return res

    good_ratio = 1.0 - (undecoded / max(1, res.n_characters))
    # Real Morse has a clean bimodal element histogram: mostly 1 and 3 units.
    units = np.array([n / dot for state, n in runs if state], dtype=np.float64)
    # Real keying lands tightly on one and three units.  A loose tolerance
    # here lets the random run lengths of a noisy channel "fit" Morse.
    fit = float(np.mean(np.minimum(np.abs(units - 1.0), np.abs(units - 3.0)) < 0.35))
    res.timing_fit = round(fit, 3)
    speed_ok = 1.0 if 5.0 <= res.wpm <= 60.0 else 0.4
    conf = 0.45 * good_ratio + 0.35 * fit + 0.20 * speed_ok
    if res.n_characters < 3:
        conf *= 0.5

    # A uniform pulse train decodes to "EEEE" or "TTTT" and looks perfect by
    # the metrics above.  Real text mixes dots and dashes and uses more than
    # one or two letters, so require both.
    dash_ratio = float(np.mean(units > 2.0)) if units.size else 0.0
    res.notes.append(f"{dash_ratio*100:.0f}% dashes")
    if dash_ratio < 0.10 or dash_ratio > 0.80:
        conf *= 0.35
        res.notes.append("element mix is not typical of Morse text")
    distinct = {p for p in text_parts if p not in (" ", "<?>")}
    if len(distinct) < 2:
        # One repeated character is only suspicious when the character is one
        # element long: "EEEE" and "TTTT" are what a uniform pulse train
        # decodes to, and they are not text.  "VVV" is text -- it is the
        # oldest thing anyone sends -- and the character has structure of its
        # own, which a train of identical pulses cannot produce.
        only = next(iter(distinct), "")
        if len(_REVERSE.get(only, "")) <= 1:
            conf *= 0.35
            res.notes.append("one repeated element, not text")
    res.confidence = round(min(0.99, max(0.0, conf)), 3)
    if res.confidence < 0.5:
        res.notes.append("timing does not fit Morse cleanly")
    return res