#!/usr/bin/env python3
"""Erzeugt die Illustrationen fuer den Stufe-3-Artikel aus echten Laeufen der
Verarbeitungskette (keine Mock-Daten)."""
import sys
sys.path.insert(0, '.')
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import pmt
from gnuradio import gr, blocks, filter as grfilter
from gnuradio.filter import firdes
from gnuradio.fft import window

from epy_noise_floor import blk as NoiseFloorBlk
from epy_element_tagger import blk as TaggerBlk
from epy_afc import blk as AfcBlk

IMG_DIR = '../../../homepage/src/assets/images'


def run_afc(wav_path, wpm, pitch, samp_rate=8000, decim=8, acq_margin=80.0):
    tb = gr.top_block()
    src = blocks.wavfile_source(wav_path, False)
    bw = max(4.0 * wpm, 100.0) + acq_margin
    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)
    nf = NoiseFloorBlk(samp_rate=working_rate, ref_wpm=30.0, num_bins=20)
    th = blocks.threshold_ff(3.0, 6.0, 0)
    tagger = TaggerBlk(samp_rate=working_rate, wpm=wpm)
    afc = AfcBlk(samp_rate=working_rate, pitch=pitch, loop_gain=0.05)
    null_sink = blocks.null_sink(gr.sizeof_float)

    tb.connect(src, xlate, c2mag, nf, th, tagger, null_sink)
    tb.connect(xlate, afc)
    tb.run()
    return working_rate, afc, tagger


def plot_afc_tracking():
    working_rate, afc, _ = run_afc('cw_test_drift.wav', wpm=20, pitch=700)
    log = afc.last_correction_log
    t = [off / working_rate for off, _ in log]
    corr = [c for _, c in log]

    # Ground Truth: linearer Drift ueber die 48 Mark-Elemente, siehe gen_cw_testsignal.py
    total_duration = t[-1] if t else 1.0
    t_true = np.linspace(0, total_duration, 100)
    true_drift = 60.0 * (t_true / total_duration)

    fig, ax = plt.subplots(figsize=(10, 4.5))
    ax.plot(t_true, true_drift, color='#4c566a', ls='--', lw=1.5, label='Tatsaechlicher Drift (Ground Truth)')
    ax.step(t, corr, color='#bf616a', lw=1.5, where='post', label='AFC-Korrektur (getrackt)')
    ax.set_xlabel('Zeit [s]')
    ax.set_ylabel('Frequenzkorrektur [Hz]')
    ax.set_title('AFC-Frequenznachfuehrung gegen 60-Hz-Drift ueber die Aufnahme (20 WPM)')
    ax.legend(loc='upper left', fontsize=9)
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(f'{IMG_DIR}/cw-grc-v3-afc-tracking.png', dpi=150)
    print("Geschrieben: cw-grc-v3-afc-tracking.png")


def plot_speed_convergence():
    # 15->30-WPM-Geschwindigkeitsrampe: zeigt, wie der PLL-Regelkreis Rate (T)
    # und konstanten Verarbeitungs-Bias (b) getrennt nachfuehrt, waehrend sich
    # die tatsaechliche Sendegeschwindigkeit kontinuierlich aendert.
    samp_rate, decim, pitch, wpm_guess = 8000, 8, 700, 15
    working_rate = samp_rate / decim
    bw = max(4.0 * wpm_guess, 100.0) + 80.0

    tb = gr.top_block()
    src = blocks.wavfile_source('cw_test_speedramp.wav', False)
    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)
    nf = NoiseFloorBlk(samp_rate=working_rate, ref_wpm=30.0, num_bins=20)
    th = blocks.threshold_ff(3.0, 6.0, 0)
    tagger = TaggerBlk(samp_rate=working_rate, wpm=wpm_guess)
    sink = blocks.vector_sink_f()

    log = []  # (sample_offset, T_mark, b_mark)
    orig_work = tagger.work

    def logging_work(input_items, output_items):
        n = orig_work(input_items, output_items)
        log.append((tagger.nitems_written(0), tagger.mark.T, tagger.mark.b))
        return n
    tagger.work = logging_work

    tb.connect(src, xlate, c2mag, nf, th, tagger, sink)
    tb.run()

    x = np.array([off for off, _, _ in log]) / working_rate
    t_y = np.array([t for _, t, _ in log]) / working_rate * 1000
    b_y = np.array([b for _, _, b in log]) / working_rate * 1000
    wpm_equiv = 1.2 / (t_y / 1000)  # WPM-Aequivalent der reinen Rate T (ohne Bias b)

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6.5), sharex=True)

    ax1.plot(x, wpm_equiv, color='#5e81ac', lw=1.6, label='WPM-Aequivalent von T_mark')
    ax1.axhline(15, color='#4c566a', ls=':', lw=1.2, label='Start (15 WPM)')
    ax1.axhline(30, color='#4c566a', ls='--', lw=1.2, label='Ende (30 WPM)')
    ax1.set_ylabel('WPM-Aequivalent (aus T_mark)')
    ax1.legend(loc='center right', fontsize=8)
    ax1.grid(True, alpha=0.3)

    ax2.plot(x, b_y, color='#bf616a', lw=1.6, label='b_mark (Verarbeitungs-Bias)')
    ax2.set_xlabel('Zeit [s]')
    ax2.set_ylabel('b_mark [ms]')
    ax2.legend(loc='center right', fontsize=8)
    ax2.grid(True, alpha=0.3)

    fig.suptitle('Geschwindigkeitsrampe 15->30 WPM: T_mark folgt der Rampe, b_mark bleibt stabil')
    fig.tight_layout()
    fig.savefig(f'{IMG_DIR}/cw-grc-v3-speed-convergence.png', dpi=150)
    print("Geschrieben: cw-grc-v3-speed-convergence.png")


if __name__ == '__main__':
    plot_afc_tracking()
    plot_speed_convergence()
