#!/usr/bin/env python3
# Two receive-path claims checked through the firmware demodulator and framer (innerpulse DM-1701 review audits):
#  1. "the 24 -> 38.4 kHz linear resampler aliases": standard C4FM direct at 38.4 kHz versus band-limited to 24 kHz and
#     then upsampled by the firmware's own p25_resample() in 32-sample chunks, as the live monitor does;
#  2. "0x34=0x00 removes the 3 kHz voice low-pass": the manual documents a choice of 2.55 or 3 kHz low-pass, not a bypass,
#     so compare no low-pass, 3 kHz and 2.55 kHz (4th-order Butterworth; the real corners and orders are undocumented).
# Run: DM1701_REPO=/path/to/DM-1701 C4FM_SIM_WORK=/path/to/work python3 resampler_check.py
import ctypes as C
import json
import math
import sys
from pathlib import Path

import numpy as np

sys.dont_write_bytecode = True

sys.path.insert(0, str(Path(__file__).resolve().parent))
import c4fm_sim as S  # noqa: E402

SNRS = (None, 16, 13, 11)
TRIALS = 3


class ResampleState(C.Structure):
    _fields_ = [('t', C.c_float), ('last', C.c_int16), ('init', C.c_int)]


def band_limit_resample(x, n_out):
    """Ideal resampling in the frequency domain; C4FM occupies < 2.88 kHz, so 38.4 -> 24 kHz loses nothing."""
    X = np.fft.rfft(x)
    Y = np.zeros(n_out // 2 + 1, dtype=complex)
    k = min(len(Y), len(X))
    Y[:k] = X[:k]
    return np.fft.irfft(Y, n_out) * (n_out / len(x))


def firmware_upsample(lib, x24):
    lib.p25_resample.argtypes = [C.POINTER(C.c_int16), C.c_int, C.c_float, C.POINTER(C.c_int16), C.c_int, C.POINTER(ResampleState)]
    lib.p25_resample.restype = C.c_int
    s = np.clip(np.round(x24), -32768, 32767).astype(np.int16)
    st = ResampleState(0.0, 0, 0)
    obuf = (C.c_int16 * 192)()
    out = []
    for off in range(0, len(s), 32):
        chunk = s[off:off + 32]
        n = lib.p25_resample((C.c_int16 * len(chunk))(*chunk.tolist()), len(chunk), C.c_float(38400.0 / 24000.0), obuf, 192, C.byref(st))
        out += list(obuf[:n])
    return np.array(out, dtype=float)


def run(lib, x, wire, wanted):
    frames, nids, dib, _ = S.decode(lib, x, len(wanted))
    errors, symbols = S.raw_ser(dib, wire)
    return frames, nids, errors, symbols


def lowpass(fc):
    return [S.biquad('lp', fc, 0.5412), S.biquad('lp', fc, 1.3066)]


def main():
    S.build_libs()
    lib = S.load('rrc')
    wire, wanted = S.build_stream(20, 1701)
    base = S.tx_c4fm(wire)
    base = base * (6000.0 / float(abs(base).max()))
    n24 = int(round(len(base) * 24000 / 38400))

    resampler = []
    for snr in SNRS:
        row = {'snr': snr, 'direct': {'frames': [], 'nids': [], 'errors': 0, 'symbols': 0},
               'via24k': {'frames': [], 'nids': [], 'errors': 0, 'symbols': 0}}
        for trial in range(TRIALS if snr is not None else 1):
            x = S.add_noise(base, snr, 777 + trial) if snr is not None else base
            for key, y in (('direct', x), ('via24k', firmware_upsample(lib, band_limit_resample(x, n24)))):
                fr, ni, e, t = run(lib, y, wire, wanted)
                row[key]['frames'].append(fr)
                row[key]['nids'].append(ni)
                row[key]['errors'] += e
                row[key]['symbols'] += t
        d, v = row['direct'], row['via24k']
        print(f"SNR {str(snr) if snr is not None else 'none':>4}: direct frames {d['frames']} errors {d['errors']}/{d['symbols']}"
              f"   via 24 kHz + p25_resample frames {v['frames']} errors {v['errors']}/{v['symbols']}", flush=True)
        resampler.append(row)

    up = firmware_upsample(lib, band_limit_resample(base, n24))
    N = 1 << 16
    P = np.abs(np.fft.rfft(up[:N] * np.hanning(N))) ** 2
    f = np.fft.rfftfreq(N, 1 / 38400.0)
    images_db = 10 * math.log10(P[(f > 3600)].sum() / P[f < 2880].sum())
    print(f'power above 3.6 kHz after p25_resample: {images_db:.1f} dB relative to the C4FM band', flush=True)

    lp = []
    for label, fc in (('no low-pass', None), ('3 kHz low-pass', 3000.0), ('2.55 kHz low-pass', 2550.0)):
        stages = lowpass(fc) if fc else []
        row = {'label': label, 'fc': fc}
        for tag, snr, trials in (('clean', None, 1), ('snr13', 13, TRIALS)):
            acc = {'frames': [], 'nids': [], 'errors': 0, 'symbols': 0}
            for trial in range(trials):
                x = S.add_noise(base, snr, 777 + trial) if snr is not None else base
                y = S.apply_chain(x, stages) if stages else x
                fr, ni, e, t = run(lib, y, wire, wanted)
                acc['frames'].append(fr)
                acc['nids'].append(ni)
                acc['errors'] += e
                acc['symbols'] += t
            row[tag] = acc
        print(f"{label:18s} noise-free frames {row['clean']['frames']} errors {row['clean']['errors']}/{row['clean']['symbols']}"
              f"   13 dB frames {row['snr13']['frames']} errors {row['snr13']['errors']}/{row['snr13']['symbols']}", flush=True)
        lp.append(row)

    out = {'stream': {'ldus': 20, 'imbe_frames': 180, 'nids': 20, 'seed': 1701, 'noise_seeds': [777 + i for i in range(TRIALS)]},
           'resampler': resampler, 'images_db': round(images_db, 1), 'lowpass': lp}
    S.WORK.mkdir(parents=True, exist_ok=True)
    (S.WORK / 'resampler_check.json').write_text(json.dumps(out, indent=1))


if __name__ == '__main__':
    main()