#!/usr/bin/env python3
"""Testet die Stufe-3-Verarbeitungskette (adaptive Geschwindigkeit, Warm-up-
Sperre, AFC) direkt in Python/GNU Radio gegen die Testsignale, bevor die
Blockverschaltung als .grc uebernommen wird."""
import sys
import numpy as np
from gnuradio import gr, blocks, filter as grfilter
from gnuradio.filter import firdes
from gnuradio.fft import window
import pmt

sys.path.insert(0, '.')
from epy_noise_floor import blk as NoiseFloorBlk
from epy_element_tagger import blk as TaggerBlk
from epy_decoder import blk as DecoderBlk
from epy_afc import blk as AfcBlk


class TextSink(gr.basic_block):
    def __init__(self):
        gr.basic_block.__init__(self, name="TextSink", in_sig=None, out_sig=None)
        self.message_port_register_in(pmt.intern("in"))
        self.set_msg_handler(pmt.intern("in"), self.handle)
        self.last = ""

    def handle(self, msg):
        self.last = pmt.to_python(msg)


def run(wav_path, wpm, pitch, samp_rate=8000, decim=8, acq_margin=80.0, verbose=False):
    tb = gr.top_block()

    src = blocks.wavfile_source(wav_path, False)

    bw = max(4.0 * wpm, 100.0) + acq_margin  # Akquisitionsbandbreite mit AFC-Reserve
    working_rate = samp_rate / decim

    xlate = grfilter.freq_xlating_fir_filter_fcf(
        decim, firdes.low_pass(1.0, samp_rate, bw / 2, bw / 2, window.WIN_HAMMING),
        pitch, samp_rate)

    c2mag = blocks.complex_to_mag(1)
    noise_floor = NoiseFloorBlk(samp_rate=working_rate, ref_wpm=30.0, num_bins=20)
    thresh = blocks.threshold_ff(3.0, 6.0, 0)
    tagger = TaggerBlk(samp_rate=working_rate, wpm=wpm)
    decoder = DecoderBlk()
    sink_text = TextSink()
    afc = AfcBlk(samp_rate=working_rate, pitch=pitch, loop_gain=0.05)

    tb.connect(src, xlate, c2mag, noise_floor, thresh, tagger, decoder)
    tb.connect(xlate, afc)
    tb.msg_connect(decoder, "out", sink_text, "in")
    tb.msg_connect(afc, "freq", xlate, "freq")

    tb.run()

    result = sink_text.last
    if verbose:
        print(f"--- {wav_path} (wpm={wpm}, pitch={pitch}, bw={bw:.1f} Hz, "
              f"working_rate={working_rate:.0f} Sa/s) ---")
        print("Decodiert:", repr(result))
        print("Finale Mark-Periode T/b:", f"{tagger.mark.T/working_rate*1000:.1f} ms / "
              f"{tagger.mark.b/working_rate*1000:.1f} ms bias")
        print("Finale Gap-Periode T/b:", f"{tagger.gap.T/working_rate*1000:.1f} ms / "
              f"{tagger.gap.b/working_rate*1000:.1f} ms bias")
        print("AFC-Korrekturverlauf (letzte 5):", afc.last_correction_log[-5:])
    return result, tagger, afc


if __name__ == '__main__':
    run('cw_test_20wpm.wav', wpm=20, pitch=700, verbose=True)
    run('cw_test_35wpm.wav', wpm=35, pitch=600, verbose=True)
    run('cw_test_20wpm_noisy.wav', wpm=20, pitch=700, verbose=True)
    run('cw_test_drift.wav', wpm=20, pitch=700, verbose=True)
    run('cw_test_speedramp.wav', wpm=15, pitch=700, verbose=True)
