/*
 * PR06 step5 / Task06-4: iq_rx_thr
 *   Receive the 10 MSPS ORx IQ window the FPGA injects into eth1 (UDP 6100) and
 *   turn each COMPLETE window into a Welch-averaged spectrum.
 *
 *   Wire format: 8-byte app header + 256 complex samples, 4 bytes each,
 *   big-endian: I[15:8] I[7:0] Q[15:8] Q[7:0].
 *     hdr[0..1] = 'I','Q'   hdr[2] = window_id   hdr[4..5] = pkt_idx (0..1023)
 *     hdr[6]    = fmt = 2 (16-bit IQ)
 *   One window = 1024 packets = 262144 samples = 26.21 ms = 1.08 ATSC fields.
 *
 *   *** Why two threads (Task06-4) ***
 *   The old single-loop version ran a 256-point FFT plus 256 cosf() calls inside
 *   the recvfrom() loop. At 5132 pkt/s the imx6 could not keep up: the socket
 *   reported drops=3810261 and only 310..930 of every 1024 packets were ever
 *   processed, so the app never saw a contiguous window and could not span an
 *   ATSC field. Now the RX thread does nothing but recvfrom + memcpy into a
 *   window buffer (~5.6 MB/s), and a worker thread computes the spectrum of a
 *   finished window out of band, double-buffered.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include <unistd.h>
#include <pthread.h>
#include <errno.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include "iq_rx.h"
#include "../inc/config.h"

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

#define IQ_NF     256      /* samples per packet */
#define IQ_NPKT   1024     /* packets per window (must match ram2_iq_ctrl) */
#define IQ_PORT   6100
#define IQ_FS_MHZ 10.0f
#define IQ_PAY    (IQ_NF * 4)          /* 1024 payload bytes per packet */
#define IQ_WBYTES (IQ_NPKT * IQ_PAY)   /* 1 MiB per window */

/* Welch over the whole window. The FFT is 4096 points because at 10 MSPS that is
 * a 2441.40625 Hz bin -- bit for bit the same bin the FPGA's imd_imd_fft_meas uses
 * at 40 MHz/16384. So bins 10 / +1290 / -1290 land exactly on its cent/L/R taps and
 * UIMD/LIMD can be recomputed here without redefining anything. The published trace
 * is folded down to 1024 points (4 bins each) to keep the JSON small. */
#define IQ_FFT    4096
#define IQ_SEGPKT (IQ_FFT / IQ_NF)     /* 16 packets per segment */
#define IQ_NSEG   (IQ_NPKT / IQ_SEGPKT)
#define IQ_OUT    1024                 /* published trace points */
#define IQ_FOLD   (IQ_FFT / IQ_OUT)

/* imd_imd_fft_meas taps: cent = +24.414 kHz, L = +3.1494 MHz, R = -3.1494 MHz.
 * HALF widens each tap to 11 bins (+-12.2 kHz): a single 2.4 kHz bin of a data
 * signal is chi-square with 2 dof, i.e. +-8 dB of scatter, and the -3.1494 MHz bin
 * additionally sits on a narrow feature ~5 dB above its neighbours. */
#define IMD_CENT  10
#define IMD_L     1290
#define IMD_R     (IQ_FFT - IMD_L)
#define IMD_HALF  5

#define IQ_JSON     "/tmp/iq_spectrum.json"
#define IQ_JSON_TMP "/tmp/iq_spectrum.json.tmp"
/* /tmp is tmpfs here, but a 1 MiB dump every 170 ms is pure waste; ~2 s is plenty
 * for the PC-side pull described in Result06-5. */
#define IQ_RAW_EVERY 12

/* The 8VSB demod is expensive (seconds, not milliseconds -- see iq_vsb.c). It runs on
 * its own low-priority thread off a private copy of the window, so a slow demod can
 * never stall the spectrum worker or the RX socket. Hand it a window only when it is
 * idle; skipped windows cost nothing. */
#define IQ_VSB_EVERY 6

/* ---- double-buffered window between the RX thread and the worker ---- */
static unsigned char  iq_win[2][IQ_WBYTES];
static unsigned char  iq_seen[2][IQ_NPKT];
static pthread_mutex_t iq_mtx  = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t  iq_cond = PTHREAD_COND_INITIALIZER;
static volatile int    iq_ready = -1;   /* buffer index handed to the worker, -1 = none */
static volatile int    iq_ready_pkts;
static volatile int    iq_ready_run;    /* longest contiguous pkt_idx run in that window */
static volatile int    iq_ready_wid;

/* Kept in RAM deliberately: every process start must default to OFF. */
static volatile int    iq_spec_enable;

/* in-place radix-2 complex FFT */
static void iq_fft(float *re, float *im, int n)
{
    int i, j, len; float ang, wr, wi, ur, ui, tr, ti;
    for (i = 1, j = 0; i < n; i++) {
        int bit = n >> 1;
        for (; j & bit; bit >>= 1) j ^= bit;
        j ^= bit;
        if (i < j) { tr = re[i]; re[i] = re[j]; re[j] = tr;
                     ti = im[i]; im[i] = im[j]; im[j] = ti; }
    }
    for (len = 2; len <= n; len <<= 1) {
        ang = -2.0f * (float)M_PI / len; wr = cosf(ang); wi = sinf(ang);
        for (i = 0; i < n; i += len) {
            ur = 1.0f; ui = 0.0f;
            for (j = 0; j < len / 2; j++) {
                int a = i + j, b = i + j + len / 2;
                tr = re[b] * ur - im[b] * ui;
                ti = re[b] * ui + im[b] * ur;
                re[b] = re[a] - tr; im[b] = im[a] - ti;
                re[a] += tr;        im[a] += ti;
                { float nur = ur * wr - ui * wi, nui = ur * wi + ui * wr; ur = nur; ui = nui; }
            }
        }
    }
}

/* longest run of consecutive packets actually received */
static int longest_run(const unsigned char *seen)
{
    int i, run = 0, best = 0;
    for (i = 0; i < IQ_NPKT; i++) {
        if (seen[i]) { run++; if (run > best) best = run; }
        else run = 0;
    }
    return best;
}

static int cmp_float(const void *a, const void *b)
{
    float x = *(const float *)a, y = *(const float *)b;
    return (x < y) ? -1 : (x > y);
}

/* ---- newest UIMD/LIMD, published to lcd_task.c's get_tsnr_limd_uimd() ---- */
static pthread_mutex_t imd_mtx = PTHREAD_MUTEX_INITIALIZER;
static float  imd_l_db, imd_u_db;
static time_t imd_stamp;            /* 0 = nothing measured yet */

int iq_get_imd(float *limd_db, float *uimd_db)
{
    int fresh;
    if (!iq_spec_get_enable()) return 0;
    pthread_mutex_lock(&imd_mtx);
    fresh = imd_stamp && (time(NULL) - imd_stamp) <= IQ_IMD_MAX_AGE;
    if (fresh) { *limd_db = imd_l_db; *uimd_db = imd_u_db; }
    pthread_mutex_unlock(&imd_mtx);
    return fresh;
}

/* ---- private hand-off to the 8VSB demod thread ---- */
static unsigned char   vsb_win[IQ_WBYTES];
static pthread_mutex_t vsb_mtx  = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t  vsb_cond = PTHREAD_COND_INITIALIZER;
static volatile int    vsb_pending;      /* 1 = vsb_win holds a window to analyse */
static volatile int    vsb_wid;

int iq_spec_get_enable(void)
{
    return iq_spec_enable != 0;
}

void iq_spec_set_enable(int enable)
{
    enable = enable ? 1 : 0;
    if (iq_spec_enable == enable) return;

    iq_spec_enable = enable;
    if (!enable) {
        pthread_mutex_lock(&iq_mtx);
        iq_ready = -1;
        pthread_mutex_unlock(&iq_mtx);

        pthread_mutex_lock(&vsb_mtx);
        vsb_pending = 0;
        pthread_mutex_unlock(&vsb_mtx);

        pthread_mutex_lock(&imd_mtx);
        imd_stamp = 0;
        pthread_mutex_unlock(&imd_mtx);
        iq_vsb_clear_results();

        unlink(IQ_JSON);
        unlink(IQ_JSON_TMP);
        unlink("/tmp/iq_eye.json");
    }
    printf("[iq_rx] spectrum %s\n", enable ? "enabled" : "disabled");
}

static void *iq_vsb_thr(void *arg)
{
    (void)arg;
#if USE_UI_RESPONSE_PRIORITY
    if (setpriority(PRIO_PROCESS, 0, 10) != 0)
        printf("[sched] iq_vsb_thr nice 10 failed: %s\n", strerror(errno));
#else
    nice(5);                             /* never outrank the RX socket or TSoIP */
#endif
    for (;;) {
        pthread_mutex_lock(&vsb_mtx);
        while (!vsb_pending) pthread_cond_wait(&vsb_cond, &vsb_mtx);
        pthread_mutex_unlock(&vsb_mtx);

        if (!iq_spec_get_enable()) continue;
        iq_vsb_analyse(vsb_win, vsb_wid);

        if (!iq_spec_get_enable()) {
            iq_vsb_clear_results();
            unlink("/tmp/iq_eye.json");
        }

        pthread_mutex_lock(&vsb_mtx);
        vsb_pending = 0;
        pthread_mutex_unlock(&vsb_mtx);
    }
    return 0;
}

/* mean power of the 2*half+1 bins centred on k, wrapping the FFT index */
static float band_power(const float *psd, int k, int half)
{
    float acc = 0.0f;
    int d;
    for (d = -half; d <= half; d++) acc += psd[(k + d) & (IQ_FFT - 1)];
    return acc / (2 * half + 1);
}

/*
 * Welch-average one window and publish it as JSON for sys_spectrum.html.
 *
 * The dB reference is the mean in-band flat power density (|f| < 2.5 MHz, which
 * excludes the -2.69 MHz pilot). Result06-5 §7 explains why: referencing the
 * pilot peak makes every number a function of the FFT length, because the pilot
 * is a CW line while the shoulder is a noise density. rel-flat is RBW-invariant.
 */
static void *iq_spec_thr(void *arg)
{
    static float hann[IQ_FFT];
    static float psd[IQ_FFT];
    static float re[IQ_FFT], im[IQ_FFT];
    static float sh[IQ_FFT];
    static char  json[24576];
    int i, s, buf, npkt, run, wid, nseg, nraw = 0, nvsb = 0;

    for (i = 0; i < IQ_FFT; i++)
        hann[i] = 0.5f * (1.0f - cosf(2.0f * (float)M_PI * i / (IQ_FFT - 1)));

    (void)arg;
#if USE_UI_RESPONSE_PRIORITY
    if (setpriority(PRIO_PROCESS, 0, 10) != 0)
        printf("[sched] iq_spec_thr nice 10 failed: %s\n", strerror(errno));
#endif
    for (;;) {
        pthread_mutex_lock(&iq_mtx);
        while (iq_ready < 0) pthread_cond_wait(&iq_cond, &iq_mtx);
        buf = iq_ready; npkt = iq_ready_pkts; run = iq_ready_run; wid = iq_ready_wid;
        pthread_mutex_unlock(&iq_mtx);

        if (!iq_spec_get_enable()) {
            pthread_mutex_lock(&iq_mtx);
            iq_ready = -1;
            pthread_mutex_unlock(&iq_mtx);
            continue;
        }

        for (i = 0; i < IQ_FFT; i++) psd[i] = 0.0f;
        nseg = 0;
        for (s = 0; s < IQ_NSEG; s++) {
            const unsigned char *base = iq_win[buf] + (size_t)s * IQ_FFT * 4;
            float mr = 0, mi = 0;
            int p, drop = 0;
            for (p = 0; p < IQ_SEGPKT; p++) if (!iq_seen[buf][s * IQ_SEGPKT + p]) drop = 1;
            if (drop) continue;                   /* never average a torn segment */
            for (i = 0; i < IQ_FFT; i++) {
                const unsigned char *q = base + i * 4;
                re[i] = (float)(short)((q[0] << 8) | q[1]);
                im[i] = (float)(short)((q[2] << 8) | q[3]);
                mr += re[i]; mi += im[i];
            }
            mr /= IQ_FFT; mi /= IQ_FFT;           /* strip the DC/LO leakage term */
            for (i = 0; i < IQ_FFT; i++) {
                re[i] = (re[i] - mr) * hann[i];
                im[i] = (im[i] - mi) * hann[i];
            }
            iq_fft(re, im, IQ_FFT);
            for (i = 0; i < IQ_FFT; i++) psd[i] += re[i] * re[i] + im[i] * im[i];
            nseg++;
        }
        if (nseg == 0) { pthread_mutex_lock(&iq_mtx); iq_ready = -1;
                         pthread_mutex_unlock(&iq_mtx); continue; }
        if (!iq_spec_get_enable()) {
            pthread_mutex_lock(&iq_mtx);
            iq_ready = -1;
            pthread_mutex_unlock(&iq_mtx);
            continue;
        }

        {
            const float df = IQ_FS_MHZ / IQ_FFT;          /* 2.44140625 kHz, == FPGA's 40MHz/16384 */
            const float dfo = df * IQ_FOLD;               /* published bin */
            float flat = 0.0f, pk = 0.0f, pilot = 0.0f, shoulder = 0.0f;
            float p_cent, p_hi, p_lo, limd, uimd;
            int nflat = 0, nsh = 0, pkbin = 0, pos = 0, j;

            for (i = 0; i < IQ_FFT; i++) {
                float f = (i < IQ_FFT / 2 ? i : i - IQ_FFT) * df;
                psd[i] /= nseg;
                if (f > -2.5f && f < 2.5f) { flat += psd[i]; nflat++; }
                if (psd[i] > pk) { pk = psd[i]; pkbin = i; }
            }
            flat = (nflat ? flat / nflat : 1e-12f) + 1e-12f;
            pilot = (pkbin < IQ_FFT / 2 ? pkbin : pkbin - IQ_FFT) * df;

            /* UIMD/LIMD on imd_imd_fft_meas's own taps, band-averaged.
             * LIMD is by definition the shoulder on the pilot's side of the carrier,
             * so key it off the measured pilot rather than off a hard-coded sign: an
             * inverted ORx spectrum then swaps the taps instead of the labels. */
            p_cent = band_power(psd, IMD_CENT, IMD_HALF);
            p_hi   = band_power(psd, IMD_L,    IMD_HALF);   /* +3.1494 MHz */
            p_lo   = band_power(psd, IMD_R,    IMD_HALF);   /* -3.1494 MHz */
            if (pilot > 0) { float t = p_hi; p_hi = p_lo; p_lo = t; }
            limd   = 10.0f * log10f(p_cent / (p_lo + 1e-12f) + 1e-12f);
            uimd   = 10.0f * log10f(p_cent / (p_hi + 1e-12f) + 1e-12f);

            pthread_mutex_lock(&imd_mtx);
            imd_l_db = limd; imd_u_db = uimd; imd_stamp = time(NULL);
            pthread_mutex_unlock(&imd_mtx);

            /* fold 4 fine bins into one published bin; N/2 is a multiple of IQ_FOLD,
             * so a group never straddles the +-Nyquist wrap */
            for (j = 0; j < IQ_OUT; j++) {
                float acc = 0.0f;
                for (i = 0; i < IQ_FOLD; i++) acc += psd[j * IQ_FOLD + i];
                sh[j] = 10.0f * log10f(acc / IQ_FOLD / flat + 1e-12f);
            }
            for (j = 0; j < IQ_OUT; j++) psd[j] = sh[j];

            for (j = 0; j < IQ_OUT; j++) {
                float f = (j < IQ_OUT / 2 ? j : j - IQ_OUT) * dfo;
                float a = f < 0 ? -f : f;
                if (a >= 3.5f && a <= 4.8f) sh[nsh++] = psd[j];   /* shoulder, clear of Nyquist */
            }
            if (nsh) {
                qsort(sh, nsh, sizeof(float), cmp_float);
                shoulder = sh[nsh / 2];
            }

            pos = snprintf(json, sizeof(json),
                "{\"wid\":%d,\"pkts\":%d,\"npkt\":%d,\"run\":%d,\"ms\":%.2f,"
                "\"nseg\":%d,\"fs\":%.3f,\"n\":%d,\"f0\":%.6f,\"df\":%.6f,"
                "\"rbw_khz\":%.2f,\"pilot\":%.4f,\"peak\":%.2f,\"shoulder\":%.2f,"
                "\"limd\":%.2f,\"uimd\":%.2f,\"imd_rbw_khz\":%.3f,\"psd\":[",
                wid, npkt, IQ_NPKT, run, run * IQ_NF / (IQ_FS_MHZ * 1e3f),
                nseg, IQ_FS_MHZ, IQ_OUT, -IQ_FS_MHZ / 2.0f + 0.5f * (IQ_FOLD - 1) * df, dfo,
                dfo * 1000.0f, pilot, 10.0f * log10f(pk / flat + 1e-12f), shoulder,
                limd, uimd, df * 1000.0f * (2 * IMD_HALF + 1));
            /* emit fftshift-ed, i.e. ascending frequency from f0 */
            for (j = 0; j < IQ_OUT && pos < (int)sizeof(json) - 16; j++) {
                int k = (j + IQ_OUT / 2) & (IQ_OUT - 1);
                pos += snprintf(json + pos, sizeof(json) - pos, "%s%.1f",
                                j ? "," : "", psd[k]);
            }
            pos += snprintf(json + pos, sizeof(json) - pos, "]}\n");

            {   FILE *f = fopen(IQ_JSON_TMP, "w");   /* write+rename: the CGI never sees a partial file */
                if (f) { fwrite(json, 1, pos, f); fclose(f); rename(IQ_JSON_TMP, IQ_JSON); }
            }
        }

        /* a full window is 26.21 ms = 1.08 ATSC fields; print only a short burst every 30s */
        {
            static time_t log_burst_start = 0;
            static int log_burst_count = 0;
            time_t now_sec = time(NULL);

            if(log_burst_start == 0 || now_sec - log_burst_start >= 30) {
                log_burst_start = now_sec;
                log_burst_count = 0;
            }
            if(log_burst_count < 3) {
                printf("[iq_rx] win %d: %d/%d pkts, longest run %d pkts = %.2f ms (%s 1 field) period\n",
                       wid, npkt, IQ_NPKT, run, run * IQ_NF / (IQ_FS_MHZ * 1e3f),
                       (run * IQ_NF / (IQ_FS_MHZ * 1e6f) > 0.0242f) ? ">" : "<");
                log_burst_count++;
            }
        }
        if (npkt == IQ_NPKT && ++nraw >= IQ_RAW_EVERY) {
            FILE *r = fopen("/tmp/iq_raw.bin", "wb");   /* whole contiguous window, for PC-side analysis */
            if (r) { fwrite(iq_win[buf], 1, IQ_WBYTES, r); fclose(r); }
            nraw = 0;
        }
        /* The demodulator needs a torn-free window: a missing packet is a hole in the
         * middle of the symbol stream and the matched filter would smear it around.
         * Copy it aside so the demod thread can take its time. */
        if (npkt == IQ_NPKT && ++nvsb >= IQ_VSB_EVERY) {
            pthread_mutex_lock(&vsb_mtx);
            if (!vsb_pending) {                 /* still chewing the last one? skip */
                memcpy(vsb_win, iq_win[buf], IQ_WBYTES);
                vsb_wid = wid; vsb_pending = 1;
                pthread_cond_signal(&vsb_cond);
                nvsb = 0;
            }
            pthread_mutex_unlock(&vsb_mtx);
        }

        pthread_mutex_lock(&iq_mtx);
        iq_ready = -1;
        pthread_mutex_unlock(&iq_mtx);
    }
    return 0;
}

void *iq_rx_thr(void *arg)
{
    int fd, n, cur = -1, buf = 0, npkt = 0, active = 0;
    struct sockaddr_in a;
    unsigned char pkt[2048];
    pthread_t spec_id;

    (void)arg;
#if USE_UI_RESPONSE_PRIORITY
    if (setpriority(PRIO_PROCESS, 0, 5) != 0)
        printf("[sched] iq_rx_thr nice 5 failed: %s\n", strerror(errno));
#endif
    /* The injected frames are addressed to 192.168.2.143, which is NOT eth1's own
     * 192.168.2.200 -- without this secondary address the kernel treats them as
     * foreign unicast and never hands them to the socket. Note busybox has no bare
     * `ip`; the absolute path is required. */
    system("/sbin/ip addr add 192.168.2.143/24 dev eth1 2>/dev/null; /sbin/ip link set eth1 up 2>/dev/null");
    system("echo 4194304 > /proc/sys/net/core/rmem_max 2>/dev/null");

    fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (fd < 0) { printf("[iq_rx] socket fail\n"); return 0; }
    memset(&a, 0, sizeof(a));
    a.sin_family = AF_INET; a.sin_port = htons(IQ_PORT); a.sin_addr.s_addr = INADDR_ANY;
    if (bind(fd, (struct sockaddr *)&a, sizeof(a)) < 0) { printf("[iq_rx] bind %d fail\n", IQ_PORT); return 0; }
    { int rb = 4 << 20; setsockopt(fd, SOL_SOCKET, SO_RCVBUF, &rb, sizeof(rb)); }

    if (pthread_create(&spec_id, NULL, iq_spec_thr, NULL) != 0) {
        printf("[iq_rx] spec thread fail\n"); return 0;
    }
    {   pthread_t vsb_id;
        if (pthread_create(&vsb_id, NULL, iq_vsb_thr, NULL) != 0)
            printf("[iq_rx] vsb thread fail\n");
        else
            printf("[iq_rx] vsb demod thread started\n");
    }
    printf("[iq_rx] listening UDP %d, %d pkts/window, spectrum off-thread\n", IQ_PORT, IQ_NPKT);

    memset(iq_seen[0], 0, IQ_NPKT);
    for (;;) {
        int wid, idx;
        if (!iq_spec_get_enable()) {
            active = 0;
            cur = -1; npkt = 0;
            memset(iq_seen[buf], 0, IQ_NPKT);
            sleep(1);
            continue;
        }
        if (!active) {
            while (recvfrom(fd, pkt, sizeof(pkt), MSG_DONTWAIT, NULL, NULL) > 0) {}
            active = 1;
        }
        n = recvfrom(fd, pkt, sizeof(pkt), 0, NULL, NULL);
        if (n < 8 + IQ_PAY) continue;
        if (pkt[0] != 0x49 || pkt[1] != 0x51 || pkt[6] != 0x02) continue;   /* 'IQ', fmt=2 */
        wid = pkt[2];
        idx = ((pkt[4] & 0x03) << 8) | pkt[5];
        if (idx >= IQ_NPKT) continue;

        if (cur < 0) cur = wid;
        if (wid != cur) {                       /* window boundary: hand this one off */
            int other = buf ^ 1;
            pthread_mutex_lock(&iq_mtx);
            if (iq_ready < 0) {                 /* worker is free -> publish, swap buffers */
                iq_ready = buf; iq_ready_pkts = npkt; iq_ready_wid = cur;
                iq_ready_run = longest_run(iq_seen[buf]);
                pthread_cond_signal(&iq_cond);
                buf = other;
            }                                   /* else: worker still busy, reuse this buffer */
            pthread_mutex_unlock(&iq_mtx);
            memset(iq_seen[buf], 0, IQ_NPKT);
            npkt = 0; cur = wid;
        }
        /* the ONLY per-packet work: one memcpy. Everything else is off-thread. */
        memcpy(iq_win[buf] + (size_t)idx * IQ_PAY, pkt + 8, IQ_PAY);
        if (!iq_seen[buf][idx]) { iq_seen[buf][idx] = 1; npkt++; }
    }
    return 0;
}