"""Generate correctly-timed Morse audio for testing (1/3/7 unit spacing)."""
import numpy as np
from bandsaunter.morse import encode_morse


def morse_keying(msg: str, wpm: float, fs: float) -> np.ndarray:
    """Key-down/key-up waveform with ITU timing: element 1, letter 3, word 7."""
    dot = 1.2 / wpm
    segs = [(0, 8)]                       # lead-in silence
    letters = encode_morse(msg).split(" ")
    for i, tok in enumerate(letters):
        if tok == "/":
            segs.append((0, 7))           # word gap
            continue
        if i > 0 and letters[i - 1] != "/":
            segs.append((0, 3))           # letter gap
        for j, el in enumerate(tok):
            if j > 0:
                segs.append((0, 1))       # element gap
            segs.append((1, 1 if el == "." else 3))
    segs.append((0, 8))
    return np.concatenate([np.full(max(1, int(u * dot * fs)), float(s))
                           for s, u in segs])


def morse_audio(msg: str, wpm: float, fs: float = 16000.0, snr_db: float = 20.0,
                tone: float = 700.0, seed: int = 0) -> np.ndarray:
    rng = np.random.default_rng(seed)
    env = morse_keying(msg, wpm, fs)
    # Soften the keying edges the way a real transmitter's shaping does.
    k = max(3, int(0.006 * fs))
    env = np.convolve(env, np.hanning(k), "same")
    env /= max(env.max(), 1e-9)
    t = np.arange(env.size) / fs
    x = env * np.sin(2 * np.pi * tone * t)
    n = np.sqrt(np.mean(x ** 2) / (10 ** (snr_db / 10.0)))
    return x + n * rng.standard_normal(x.size)