#!/usr/bin/env python3
"""Erzeugt synthetische CW-Testsignale (WAV, mono) mit bekanntem Klartext.

Generisch nutzbar fuer beliebige WPM/Pitch/Rauschpegel/Flankenform-Kombinationen,
nicht auf eine feste Testaufnahme zugeschnitten. Dient der Verifikation des
GNU-Radio-Flowgraphs (cw.grc) gegen einen bekannten Ground-Truth-Text.

Stufe 3 erweitert dies um lineare Frequenzdrift (--drift-hz) und
Geschwindigkeitswechsel (--wpm-end) ueber die Dauer der Aufnahme, um AFC und
adaptive Geschwindigkeitsschaetzung gegen realistische, sich veraendernde
Signale zu testen statt nur gegen stationaere Parameter.
"""
import numpy as np
from scipy.io import wavfile
import argparse

MORSE = {
    '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': '----.',
}


def shaped_envelope(n_samples, rise_samples):
    """Raised-Cosine-Flanke (0..1), rise_samples lang; Plateau bleibt bei 1."""
    env = np.ones(n_samples)
    r = min(rise_samples, n_samples // 2)
    if r > 0:
        ramp = 0.5 * (1 - np.cos(np.pi * np.arange(r) / r))
        env[:r] = ramp
        env[-r:] = ramp[::-1]
    return env


def text_to_elements(text):
    """Wandelt Text in eine Liste von ('mark'|'igap'|'cgap'|'wgap', dauer_in_dots)."""
    elements = []
    words = text.upper().split(' ')
    for wi, word in enumerate(words):
        for ci, char in enumerate(word):
            code = MORSE.get(char)
            if code is None:
                continue
            for si, sym in enumerate(code):
                elements.append(('mark', 3 if sym == '-' else 1))
                if si < len(code) - 1:
                    elements.append(('igap', 1))
            if ci < len(word) - 1:
                elements.append(('cgap', 3))
        if wi < len(words) - 1:
            elements.append(('wgap', 7))
    return elements


def generate(text, wpm, pitch, fs, rise_ms, snr_db, out_path,
             drift_hz=0.0, wpm_end=None):
    rise_samples = int(fs * rise_ms / 1000)
    elements = text_to_elements(text)
    n_elem = max(1, len(elements))

    # Vorlaufzeit ausreichend lang fuer das Einschwingen des Rauschflur-Schaetzers
    # (siehe Artikeltext: ca. 5 Bin-Perioden = 5 Dot-Einheiten Einschwingzeit)
    lead_units = 8
    dot_s0 = 1.2 / wpm
    chunks = [np.zeros(int(fs * dot_s0 * lead_units))]
    t_marks = []  # (start_s, end_s) je Mark-Element, fuer die Ground-Truth-Pruefung
    pitch_track = []  # (t_s, pitch_hz) je Mark-Element, fuer die AFC-Auswertung

    cursor_s = dot_s0 * lead_units
    for i, (kind, units) in enumerate(elements):
        progress = i / (n_elem - 1) if n_elem > 1 else 0.0
        local_wpm = wpm + (wpm_end - wpm) * progress if wpm_end else wpm
        local_pitch = pitch + drift_hz * progress
        dot_s = 1.2 / local_wpm

        n = int(fs * dot_s * units)
        if kind == 'mark':
            tt = np.arange(n) / fs
            tone = np.sin(2 * np.pi * local_pitch * tt)
            tone *= shaped_envelope(n, rise_samples)
            chunks.append(tone)
            t_marks.append((cursor_s, cursor_s + n / fs))
            pitch_track.append((cursor_s, local_pitch))
        else:
            chunks.append(np.zeros(n))
        cursor_s += n / fs

    # Kurzer Abschluss-Dit nach Wortpause, damit der letzte Buchstabe noch eine
    # abschliessende steigende Flanke bekommt (Tagger deklariert ein Zeichen erst
    # bei der naechsten Flanke fuer abgeschlossen - im Dauerbetrieb kein Thema,
    # da dort immer weiteres Material folgt).
    end_pitch = pitch + drift_hz
    end_dot_s = 1.2 / (wpm_end or wpm)
    n_end = int(fs * end_dot_s)
    tt = np.arange(n_end) / fs
    end_tone = np.sin(2 * np.pi * end_pitch * tt) * shaped_envelope(n_end, rise_samples)
    chunks.append(np.zeros(int(fs * end_dot_s * 7)))
    chunks.append(end_tone)
    chunks.append(np.zeros(int(fs * end_dot_s * 2)))
    sig = np.concatenate(chunks).astype(np.float64)

    if snr_db is not None:
        sig_power = np.mean(sig[sig != 0] ** 2) if np.any(sig != 0) else 1.0
        noise_power = sig_power / (10 ** (snr_db / 10))
        rng = np.random.default_rng(42)
        sig = sig + rng.normal(0, np.sqrt(noise_power), size=sig.shape)

    sig = sig / max(np.max(np.abs(sig)), 1e-9) * 0.9
    wavfile.write(out_path, fs, sig.astype(np.float32))
    return t_marks, pitch_track, dot_s0


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('--text', default='CQ CQ DE DL1ABC TEST')
    ap.add_argument('--wpm', type=float, default=20)
    ap.add_argument('--wpm-end', type=float, default=None,
                     help='Wenn gesetzt: WPM rampt linear von --wpm zu --wpm-end ueber die Aufnahme.')
    ap.add_argument('--pitch', type=float, default=700)
    ap.add_argument('--drift-hz', type=float, default=0.0,
                     help='Lineare Tonhoehendrift ueber die Aufnahme (Endwert = pitch + drift_hz).')
    ap.add_argument('--fs', type=int, default=8000)
    ap.add_argument('--rise-ms', type=float, default=5)
    ap.add_argument('--snr-db', type=float, default=40.0,
                     help='Realistisches Signal hat immer einen Rauschflur ungleich '
                          'Null; auch der "saubere" Default-Fall nutzt daher 40 dB SNR '
                          'statt exakt Null Rauschen.')
    ap.add_argument('--out', default='cw_test.wav')
    args = ap.parse_args()

    marks, pitch_track, dot_s = generate(
        args.text, args.wpm, args.pitch, args.fs, args.rise_ms, args.snr_db,
        args.out, drift_hz=args.drift_hz, wpm_end=args.wpm_end)
    print(f"Geschrieben: {args.out} ({len(marks)} Mark-Elemente, "
          f"Dot-Laenge {dot_s*1000:.1f} ms, WPM={args.wpm}->{args.wpm_end or args.wpm}, "
          f"Pitch={args.pitch}->{args.pitch+args.drift_hz} Hz, SNR={args.snr_db})")
