#!/usr/bin/env python3
# Firmware-in-the-loop C4FM receive model from the innerpulse review audits of the DM-1701 P25 firmware.
# Run against a copy of the firmware tree:  DM1701_REPO=/path/to/DM-1701 C4FM_SIM_WORK=/path/to/work python3 c4fm_sim.py
# Needs Python 3.10+, numpy and gcc. Builds and results go to C4FM_SIM_WORK (default: <tmp>/c4fm_sim), never into the tree.
"""C4FM receive-path model driven through the firmware's own demodulator.

Question: how do (a) the RRC receive filter and (b) analog-FM-mode audio
processing affect a *standard* P25 C4FM waveform, compared with the
project's self-test model (RRC transmit -> RRC receive)?

Everything that decides dibits is the real firmware code
(p25_4fsk_demodulate_stream with the live 768/512 window, p25rx framing,
BCH NID, IMBE framing), compiled for the host exactly like tests/p25/run.py.

Waveform models (all at the demodulator's 38.4 kHz / 8 samples-per-symbol
domain, resampler not modelled):
  project  : firmware p25_channel_generate()  (RRC transmit pulse)
  c4fm     : TIA-102.BAAA C4FM: Nyquist raised cosine (alpha 0.2, 1920/2880 Hz)
             x shaping filter P(f) = (pi f/4800)/sin(pi f/4800), |f| < 2880 Hz
Receive-path models (ASSUMED corner frequencies; the AT1846S/HR-C6000
responses are undocumented, so these are illustrative, not measured):
  hpf150/hpf300 : 2nd-order Butterworth high-pass
  lpf3k         : 4th-order Butterworth low-pass at 3 kHz ("3KHz audio filter")
  deemph        : 1st-order low-pass at 300 Hz (-6 dB/oct voice de-emphasis)
Receive filters: firmware 81-tap RRC (as shipped) or an 8-sample
integrate-and-dump (TIA reference receiver), built as a scratch copy.
Noise: optional white Gaussian noise scaled so in-band (0-2.4 kHz) noise
RMS is 600 Hz * 10^(-SNR/20).  Simplification: a real FM discriminator's
noise spectrum rises with frequency.
"""
import ctypes as C
import json
import os
import subprocess
import math
import random
import re
import sys
import tempfile
from pathlib import Path

import numpy as np

sys.dont_write_bytecode = True  # importing tests/p25/run.py must not write __pycache__ into the firmware tree

HERE = Path(__file__).resolve().parent
REPO = Path(os.environ.get('DM1701_REPO', '.')).resolve()
WORK = Path(os.environ.get('C4FM_SIM_WORK', Path(tempfile.gettempdir()) / 'c4fm_sim'))
sys.path.insert(0, str(REPO / 'tests/p25'))
import run as p25run  # noqa: E402  (framing helpers: sync, independent BCH NID, 864-dibit LDUs)

FS = 38400.0
SPS = 8
LEVEL = {0: 1.0, 1: 3.0, 2: -1.0, 3: -3.0}
FW_RRC = None  # float RRC taps parsed from firmware source


LIVE_REL = 'opengd77-rt3s-experiments/MDUV380_firmware/application/source/p25'
DEMO_REL = 'dm1701-p25-demo/src/p25'


def build_libs():
    """Compile the firmware P25 sources for the host exactly as tests/p25/run.py does, three times:
    shipped RRC, an 8-sample integrate-and-dump, and a band-limited integrate-and-dump FIR (81 taps, Q12)."""
    WORK.mkdir(parents=True, exist_ok=True)
    live, demo = REPO / LIVE_REL, REPO / DEMO_REL
    fsk = (live / 'p25_4fsk.c').read_text()
    m = re.search(r'static const int16_t p25_rrc_q12\[RRC_TAPS\] = \{(.*?)\};', fsk, re.S)
    dc = sum(int(x) for x in m.group(1).replace('\n', ' ').split(',') if x.strip())
    box = [0] * 81
    for k in range(36, 44):
        box[k] = round(dc / 8)
    M = 1 << 15
    f = np.fft.fftfreq(M, 1 / FS)
    af = np.abs(f)
    hd = np.where(af <= 2880, np.sinc(af / 4800), 0.0)
    t = (af > 2880) & (af < 3600)
    hd[t] = np.sinc(2880 / 4800) * 0.5 * (1 + np.cos(np.pi * (af[t] - 2880) / 720))
    h = np.fft.fftshift(np.real(np.fft.ifft(hd)))[M // 2 - 40:M // 2 + 41] * np.kaiser(81, 5.0)
    h /= h.sum()
    q = [int(v) for v in np.round(h * 4096)]
    q[40] += 4096 - sum(q)
    (WORK / 'bldump_taps.json').write_text(json.dumps(q))
    variants = {'rrc': None, 'boxcar': box, 'bldump': q}
    for name, taps in variants.items():
        dem = live / 'p25_4fsk.c'
        if taps is not None:
            dem = WORK / f'p25_4fsk_{name}.c'
            dem.write_text(fsk[:m.start()] + 'static const int16_t p25_rrc_q12[RRC_TAPS] = {' + ','.join(map(str, taps)) + '};' + fsk[m.end():])
        cmd = ['gcc', '-O2', f'-I{demo}', f'-I{live}', f'-I{REPO / "ref-sources/dsd-master/include"}', '-DP25_DEMOD_WINDOW=768',
               '-shared', '-fPIC', str(REPO / 'tests/p25/host.c'), str(REPO / 'tests/p25/trunk_host.c')]
        cmd += [str(live / n) for n in ('p25rx.c', 'p25_nid.c', 'p25_fec.c', 'p25_trunk.c')] + [str(dem)]
        cmd += [str(live / 'p25mon_capture.c'), str(live / 'p25_resample.c')] + [str(p) for p in sorted((demo / 'mbelib').glob('*.c'))]
        cmd += ['-lm', '-o', str(WORK / f'libp25_{name}.so')]
        subprocess.run(cmd, check=True, capture_output=True)


def firmware_taps():
    src = (REPO / 'opengd77-rt3s-experiments/MDUV380_firmware/application/source/p25/p25_4fsk.c').read_text()
    import re
    fl = re.search(r'static const float p25_rrc\[RRC_TAPS\] = \{(.*?)\};', src, re.S).group(1)
    q12 = re.search(r'static const int16_t p25_rrc_q12\[RRC_TAPS\] = \{(.*?)\};', src, re.S).group(1)
    f = [float(x.strip().rstrip('f')) for x in fl.replace('\n', ' ').split(',') if x.strip()]
    q = [int(x) for x in q12.replace('\n', ' ').split(',') if x.strip()]
    return np.array(f), np.array(q, dtype=float)


def load(variant):
    lib = p25run.setup_lib(WORK / f'libp25_{variant}.so')
    lib.p25_channel_generate.argtypes = [C.POINTER(C.c_uint8), C.c_int, C.c_int, C.c_int, C.c_int,
                                         C.POINTER(C.c_int16), C.c_int]
    lib.p25_4fsk_demodulate_stream.argtypes = [C.POINTER(C.c_int16), C.c_int, C.POINTER(C.c_uint8),
                                               C.c_int, C.c_int, C.POINTER(p25run.Timing)]
    return lib


def build_stream(n_ldu, seed):
    p25run.RNG.seed(seed)
    wire, wanted = [0] * 40, []
    for k in range(n_ldu):
        frame, payload = p25run.ldu(5 if k % 2 == 0 else 10)
        wire += frame
        wanted += payload
    wire += [0] * 80
    return wire, wanted


# ---------------- transmit models ----------------
def tx_c4fm(wire):
    n = len(wire) * SPS
    pad = 2048
    N = 1 << int(math.ceil(math.log2(n + 2 * pad)))
    x = np.zeros(N)
    for i, d in enumerate(wire):
        x[pad + i * SPS] = LEVEL[d] * 600.0 * SPS
    f = np.fft.rfftfreq(N, 1.0 / FS)
    H = np.where(f <= 1920, 1.0, np.where(f <= 2880, 0.5 + 0.5 * np.cos(2 * np.pi * f / 1920.0), 0.0))
    P = np.ones_like(f)
    nz = (f > 0) & (f < 2880)
    P[nz] = (np.pi * f[nz] / 4800.0) / np.sin(np.pi * f[nz] / 4800.0)
    y = np.fft.irfft(np.fft.rfft(x) * H * P, N)
    return y[pad:pad + n]


def tx_project(lib, wire):
    src = (C.c_uint8 * len(wire))(*wire)
    out = (C.c_int16 * (len(wire) * SPS))()
    n = lib.p25_channel_generate(src, len(wire), 99, 0, 0, out, len(out))
    return np.array(out[:n], dtype=float)


# ---------------- receive-path IIR models ----------------
def lfilter(b, a, x):
    y = np.zeros_like(x)
    x1 = x2 = y1 = y2 = 0.0
    b0, b1, b2 = b
    _, a1, a2 = a
    for i, xi in enumerate(x):
        yi = b0 * xi + b1 * x1 + b2 * x2 - a1 * y1 - a2 * y2
        x2, x1 = x1, xi
        y2, y1 = y1, yi
        y[i] = yi
    return y


def biquad(kind, fc, q):
    K = math.tan(math.pi * fc / FS)
    norm = 1.0 / (1.0 + K / q + K * K)
    a1 = 2.0 * (K * K - 1.0) * norm
    a2 = (1.0 - K / q + K * K) * norm
    if kind == 'lp':
        b0 = K * K * norm
        return (b0, 2 * b0, b0), (1.0, a1, a2)
    b0 = norm
    return (b0, -2 * b0, b0), (1.0, a1, a2)


def onepole_lp(fc):
    K = math.tan(math.pi * fc / FS)
    b0 = K / (K + 1.0)
    return (b0, b0, 0.0), (1.0, (K - 1.0) / (K + 1.0), 0.0)


def rx_chain(name):
    stages = []
    for part in name.split('+'):
        if part.startswith("hpf") and float(part[3:]) > 0:
            stages.append(biquad('hp', float(part[3:]), 1 / math.sqrt(2)))
        elif part == 'lpf3k':
            stages.append(biquad('lp', 3000.0, 0.5412))
            stages.append(biquad('lp', 3000.0, 1.3066))
        elif part == 'deemph':
            stages.append(onepole_lp(300.0))
    return stages


def apply_chain(x, stages):
    for b, a in stages:
        x = lfilter(b, a, x)
    return x


def response_db(stages, freqs):
    out = []
    for f in freqs:
        z = np.exp(-1j * 2 * np.pi * f / FS)
        h = 1.0 + 0j
        for b, a in stages:
            h *= (b[0] + b[1] * z + b[2] * z * z) / (a[0] + a[1] * z + a[2] * z * z)
        out.append(20 * math.log10(max(abs(h), 1e-9)))
    return out


def add_noise(x, snr_db, seed):
    if snr_db is None:
        return x
    rng = np.random.default_rng(seed)
    unit = float(np.sqrt(np.mean(x * x)) / math.sqrt(5.0))   # scale-free: RMS of a 4-level signal = sqrt(5) x unit
    sigma_inband = unit * 10 ** (-snr_db / 20.0)
    sigma = sigma_inband * math.sqrt((FS / 2) / 2400.0)
    return x + rng.normal(0.0, sigma, size=x.shape)


# ---------------- firmware decode ----------------
def decode(lib, samples, wanted_len):
    s = np.clip(np.round(samples), -32768, 32767).astype(np.int16)
    buf = (C.c_int16 * len(s))(*s.tolist())
    out = (C.c_uint8 * 104)()
    timing = p25run.Timing()
    lib.test_init()
    dibits = []
    for off in range(0, len(s) - 767, 512):
        win = C.cast(C.byref(buf, off * 2), C.POINTER(C.c_int16))
        n = lib.p25_4fsk_demodulate_stream(win, 768, out, 104, 512, C.byref(timing))
        lib.test_feed(out, n)
        dibits += list(out[:n])
    frames = lib.test_metric(0)
    nids = lib.test_metric(1)
    got = list(lib.test_payload()[:lib.test_payload_len()])
    return frames, nids, dibits, got


def raw_ser(dibits, wire):
    """Symbol error rate of the demodulated stream against the transmitted wire,
    re-aligned every 400 symbols (tolerates timing slips)."""
    errors = total = 0
    base = None
    pos = 0
    block = 400
    while pos + block <= len(dibits):
        seg = dibits[pos:pos + block]
        best = None
        centre = base + pos if base is not None else None
        rng_ = range(centre - 4, centre + 5) if centre is not None else range(0, 400)
        for off in rng_:
            if off < 0 or off + block > len(wire):
                continue
            e = sum(1 for a, b in zip(seg, wire[off:off + block]) if a != b)
            if best is None or e < best[0]:
                best = (e, off)
        if best is None:
            break
        base = best[1] - pos
        errors += best[0]
        total += block
        pos += block
    return errors, total


# ---------------- eye diagrams (python replica of the firmware MF) ----------------
def eye(samples, wire, rxfilter, q12, n_traces=160, up=4):
    if rxfilter == 'rrc':
        taps = q12 / 4096.0
        centre = 40.0
    elif rxfilter == 'bldump':
        taps = np.array(json.load(open(WORK / 'bldump_taps.json')), dtype=float) / 4096.0
        centre = 40.0
    else:
        taps = np.zeros(81)
        taps[36:44] = 1.0 / 8
        centre = 39.5
    mf = np.convolve(samples, taps[::-1], mode='valid')          # mf[j] ~ input j + centre
    levels = np.array([LEVEL[d] for d in wire])
    first, last = 60, len(wire) - 60
    best = None
    for shift in np.arange(-8.0, 40.0, 0.25):   # IIR models add up to ~25 samples of group delay
        idx = np.arange(first, last) * SPS - centre + shift
        v = np.interp(idx, np.arange(len(mf)), mf)
        scale = math.sqrt(5.0) / math.sqrt(np.mean(v * v))
        v = v * scale
        lv = levels[first:last]
        gaps = []
        for thr_lo, thr_hi in ((-3, -1), (-1, 1), (1, 3)):
            lo = v[lv == thr_lo]
            hi = v[lv == thr_hi]
            gaps.append((hi.min() - lo.max()) / 2.0)
        opening = min(gaps)
        if best is None or opening > best[0]:
            best = (opening, shift, scale)
    opening, shift, scale = best
    # ideal-timing slicer SER (independent of the firmware's timing/level loops)
    idx = np.arange(first, last) * SPS - centre + shift
    v = np.interp(idx, np.arange(len(mf)), mf) * scale
    lv = levels[first:last]
    sliced = np.where(v >= 2, 3.0, np.where(v >= 0, 1.0, np.where(v >= -2, -1.0, -3.0)))
    ideal_ser = float(np.mean(sliced != lv))
    rng = random.Random(7)
    traces = []
    fine = np.arange(-SPS, SPS + 0.001, 1.0 / up)
    for _ in range(n_traces):
        k = rng.randrange(first + 2, last - 2)
        idx = k * SPS - centre + shift + fine
        v = np.interp(idx, np.arange(len(mf)), mf) * scale
        traces.append([round(float(t), 3) for t in v])
    return {'opening': round(float(opening), 3), 'shift': float(shift), 'ideal_ser': round(ideal_ser, 4),
            'traces': traces}


def main():
    build_libs()
    f_taps, q12 = firmware_taps()
    libs = {'rrc': load('rrc'), 'boxcar': load('boxcar'), 'bldump': load('bldump')}
    wire, wanted = build_stream(20, 1701)
    n_expected = 9 * 20
    scenarios = [
        # (key, label, tx, rx chain, rx filter)
        ('S0', 'Project test model: RRC transmit -> firmware RRC', 'project', '', 'rrc'),
        ('S1', 'Standard C4FM -> firmware RRC (as shipped)', 'c4fm', '', 'rrc'),
        ('S2', 'Standard C4FM -> integrate-and-dump (TIA reference)', 'c4fm', '', 'boxcar'),
        ('S2b', 'Standard C4FM -> band-limited integrate-and-dump FIR (81 taps)', 'c4fm', '', 'bldump'),
        ('S3', 'C4FM + 150 Hz high-pass -> RRC', 'c4fm', 'hpf150', 'rrc'),
        ('S4', 'C4FM + 300 Hz high-pass -> RRC', 'c4fm', 'hpf300', 'rrc'),
        ('S5', 'C4FM + 300 Hz HPF + 3 kHz LPF -> RRC', 'c4fm', 'hpf300+lpf3k', 'rrc'),
        ('S6', 'C4FM + 300 Hz HPF + 3 kHz LPF + de-emphasis -> RRC', 'c4fm', 'hpf300+lpf3k+deemph', 'rrc'),
        ('S7', 'C4FM + 300 Hz HPF + 3 kHz LPF + de-emphasis -> integrate-and-dump', 'c4fm', 'hpf300+lpf3k+deemph', 'boxcar'),
    ]
    base_c4fm = tx_c4fm(wire)
    base_proj = tx_project(libs['rrc'], wire)
    results = {'expected_frames': n_expected, 'expected_nids': 20, 'scenarios': []}
    for key, label, tx, chain, rxf in scenarios:
        base = base_proj if tx == 'project' else base_c4fm          # discriminator output, Hz units
        stages = rx_chain(chain) if chain else []
        clean = apply_chain(base, stages) if stages else base
        row = {'key': key, 'label': label, 'tx': tx, 'chain': chain or 'none', 'rxfilter': rxf, 'runs': {}}
        for snr in (None, 20, 14):
            # noise enters at the discriminator, BEFORE the audio-path filters
            y = add_noise(base, snr, 1234 + (snr or 0))
            y = apply_chain(y, stages) if stages else y
            # common int16 scaling (the demodulator's level fit normalises amplitude)
            y = y * (6000.0 / max(1.0, float(np.max(np.abs(clean)))))
            frames, nids, dibits, got = decode(libs[rxf], y, len(wanted))
            err, tot = raw_ser(dibits, wire)
            payload_err = sum(a != b for a, b in zip(got, wanted)) if frames == n_expected else None
            row['runs']['clean' if snr is None else f'snr{snr}'] = {
                'imbe_frames': frames, 'nids_valid': nids,
                'raw_symbol_errors': err, 'raw_symbols': tot,
                'ser': (err / tot) if tot else None,
                'payload_dibit_errors': payload_err,
            }
            print(f'{key} {("clean" if snr is None else "snr%d" % snr):>6}: frames {frames:3d}/{n_expected} '
                  f'nids {nids:2d}/20 raw SER {err}/{tot} = {err / max(tot, 1):.4f}', flush=True)
        row['eye'] = eye(clean, wire, rxf, q12)
        print(f'{key} eye opening (worst of 3, 1.0 = ideal): {row["eye"]["opening"]}  '
              f'ideal-timing SER {row["eye"]["ideal_ser"]}', flush=True)
        results['scenarios'].append(row)

    # SNR sweep: margin of the shipped RRC receiver on the project model vs standard C4FM
    sweep = {'snrs': [6, 8, 10, 12, 14, 16, 20], 'series': {}}
    for name, base, rxf in (('project_rrc', base_proj, 'rrc'), ('c4fm_rrc', base_c4fm, 'rrc'),
                            ('c4fm_intdump', base_c4fm, 'boxcar'), ('c4fm_bldump', base_c4fm, 'bldump')):
        pts = []
        for snr in sweep['snrs']:
            errs = tots = nids = frames = 0
            for seed in range(3):
                y = add_noise(base, snr, 9000 + 97 * seed + snr)
                y = y * (6000.0 / max(1.0, float(np.max(np.abs(base)))))
                fr, ni, dibits, _ = decode(libs[rxf], y, len(wanted))
                e, t = raw_ser(dibits, wire)
                errs += e; tots += t; nids += ni; frames += fr
            pts.append({'snr': snr, 'ser': errs / max(tots, 1), 'nids': nids / 3.0, 'frames': frames / 3.0})
            print(f'sweep {name:13s} snr {snr:2d}: SER {errs / max(tots, 1):.4f} nids {nids / 3:.1f} frames {frames / 3:.1f}', flush=True)
        sweep['series'][name] = pts
    results['snr_sweep'] = sweep

    # High-pass corner sweep (2nd-order Butterworth), standard C4FM, shipped RRC receiver
    hp = {'corners': [0, 10, 20, 35, 50, 75, 100, 150, 200, 300], 'clean': [], 'snr20': []}
    for fc in hp['corners']:
        stages = rx_chain(f'hpf{fc}') if fc else []
        for tag, snr in (('clean', None), ('snr20', 20)):
            y = add_noise(base_c4fm, snr, 4242 + fc)
            y = apply_chain(y, stages) if stages else y
            ref = apply_chain(base_c4fm, stages) if stages else base_c4fm
            y = y * (6000.0 / max(1.0, float(np.max(np.abs(ref)))))
            fr, ni, dibits, _ = decode(libs['rrc'], y, len(wanted))
            e, t = raw_ser(dibits, wire)
            hp[tag].append({'fc': fc, 'ser': e / max(t, 1), 'nids': ni, 'frames': fr})
            print(f'hpf {fc:3d} Hz {tag:5s}: SER {e / max(t, 1):.4f} nids {ni}/20 frames {fr}/180', flush=True)
    results['hpf_sweep'] = hp

    # frequency responses for plotting (dB)
    freqs = [float(f) for f in np.round(np.geomspace(50, 12000, 120), 1)]
    f = np.array(freqs)
    H = np.where(f <= 1920, 1.0, np.where(f <= 2880, 0.5 + 0.5 * np.cos(2 * np.pi * f / 1920.0), 1e-6))
    P = np.where(f < 2880, (np.pi * f / 4800.0) / np.sin(np.pi * f / 4800.0), 1.0)
    D = np.abs(np.sinc(f / 4800.0))
    def fir_db(taps):
        w = np.exp(-1j * 2 * np.pi * np.outer(f, np.arange(len(taps))) / FS)
        h = np.abs(w @ taps)
        return [round(float(20 * np.log10(max(v, 1e-6))), 2) for v in h / (abs(np.sum(taps)))]
    results['responses'] = {
        'freqs': freqs,
        'c4fm_tx': [round(float(20 * np.log10(max(v, 1e-6))), 2) for v in H * P],
        'rrc_fw': fir_db(q12),
        'bldump_fir': fir_db(np.array(json.load(open(WORK / 'bldump_taps.json')), dtype=float)),
        'integrate_dump': [round(float(20 * np.log10(max(v, 1e-6))), 2) for v in D],
        'hpf300': [round(v, 2) for v in response_db(rx_chain('hpf300'), freqs)],
        'lpf3k': [round(v, 2) for v in response_db(rx_chain('lpf3k'), freqs)],
        'deemph': [round(v, 2) for v in response_db(rx_chain('deemph'), freqs)],
        'analog_total': [round(v, 2) for v in response_db(rx_chain('hpf300+lpf3k+deemph'), freqs)],
    }
    (WORK / 'results.json').write_text(json.dumps(results))
    print('wrote', WORK / 'results.json')


if __name__ == '__main__':
    main()