/*---------------------------------------------------------------------------*\

  FILE........: ch.c
  AUTHOR......: David Rowe
  DATE CREATED: May 2015

  Channel simulation program for testing command line versions of modems.

\*---------------------------------------------------------------------------*/

/*
  Copyright (C) 2015-2022 David Rowe

  All rights reserved.

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU Lesser General Public License version 2.1, as
  published by the Free Software Foundation.  This program is
  distributed in the hope that it will be useful, but WITHOUT ANY
  WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
  License for more details.

  You should have received a copy of the GNU Lesser General Public License
  along with this program; if not, see <http://www.gnu.org/licenses/>.
*/

#include <assert.h>
#include <errno.h>
#include <getopt.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "codec2_cohpsk.h"
#include "comp_prim.h"
#include "debug_alloc.h"
#include "freedv_api.h"
#include "ht_coeff.h"
#include "ssbfilt_coeff.h"

#define BUF_N 160
#define MPG_DELAY_MS 0.5
#define MPP_DELAY_MS 2.0
#define MPD_DELAY_MS 4.0

/* see instructions below for how to generate these files */

#define DEFAULT_FADING_DIR "unittest"
#define MPG_FADING_FILE_NAME "slow_fading_samples.float"
#define MPP_FADING_FILE_NAME "fast_fading_samples.float"
#define MPD_FADING_FILE_NAME "faster_fading_samples.float"

// Gaussian from uniform:
float gaussian(void) {
  double x = (double)rand() / RAND_MAX;
  double y = (double)rand() / RAND_MAX;
  double z = sqrt(-2 * log(x)) * cos(2 * M_PI * y);
  return sqrt(1. / 2.) * z;
}

// complex noise sample
COMP noise(void) {
  COMP n = {gaussian(), gaussian()};
  return n;
}

int main(int argc, char *argv[]) {
  FILE *fin, *ffading, *fout;
  char *fading_dir;
  float NodB, foff_hz;
  int fading_en, nhfdelay;

  short buf[BUF_N];
  float htbuf[HT_N + BUF_N];
  COMP ch_in[BUF_N];
  COMP ch_fdm[BUF_N];
  COMP ssbfiltbuf[SSBFILT_N + BUF_N];
  COMP ssbfiltout[BUF_N];

  COMP phase_ch;
  float No, variance;
  COMP scaled_noise;
  float hf_gain;
  COMP *ch_fdm_delay = NULL, aspread, aspread_2ms, delayed, direct;
  float tx_pwr, tx_pwr_fade, noise_pwr, user_multipath_delay;
  int frames, i, j, k, Fs, ret, nclipped, noutclipped, ssbfilt_en, complex_out,
      ctest;
  float sam, peak, clip, papr, CNo, snr3k, gain;

  if (argc < 3) {
  helpmsg:
    fprintf(
        stderr,
        "Command line channel simulation tool.\n"
        "\n"
        "usage: %s InputRealModemRawFile OutputRealModemRawFile [Options]\n"
        "\n"
        "  real int16 input -> Gain -> Hilbert Transform -> clipper -> freq "
        "shift ->\n"
        "  Multipath -> AWGN noise -> SSB filter -> real int16 output\n"
        "\n"
        "[--clip int16]         Hilbert clipper (clip complex signal "
        "magnitude, default 32767)\n"
        "[--complexout]         Optional int16 IQ complex output (default real "
        "int16)\n"
        "[--ctest]              Check PAPR is around 0dB, used to support "
        "ctests\n"
        "[--freq FoffHz]        Frequency offset (default 0Hz)\n"
        "[--fading_dir Path]    path to multipath fading files (default "
        "'unittest')\n"
        "[--Fs SampleRateHz]    Sample rate of simulation (default 8000 Hz)\n"
        "[--gain G]             Linear gain (default 1.0)\n"
        "[--mpg]                Multipath good 0.1Hz Doppler, 0.5ms delay\n"
        "[--mpp]                Multipath poor 1.0Hz Doppler, 2.0ms delay\n"
        "[--mpd]                Multipath disturbed 2.0Hz Doppler, 4.0ms "
        "delay\n"
        "[--ssbfilt 0|1]        SSB bandwidth filter (default 1 on)\n"
        "[--mulipath_delay ms]  Optionally adjust multipath delay\n"
        "[--No dBHz]            AWGN Noise density dB/Hz (default -100)"
        "\n",
        argv[0]);
    exit(1);
  }

  if (strcmp(argv[1], "-") == 0)
    fin = stdin;
  else if ((fin = fopen(argv[1], "rb")) == NULL) {
    fprintf(stderr, "ch: Error opening input modem raw file: %s: %s.\n",
            argv[1], strerror(errno));
    exit(1);
  }

  if (strcmp(argv[2], "-") == 0)
    fout = stdout;
  else if ((fout = fopen(argv[2], "wb")) == NULL) {
    fprintf(stderr, "ch: Error opening output modem raw file: %s: %s.\n",
            argv[2], strerror(errno));
    exit(1);
  }

  NodB = -100;
  Fs = 8000;
  foff_hz = 0.0;
  fading_en = 0;
  ctest = 0;
  clip = 32767;
  gain = 1.0;
  ssbfilt_en = 1;
  complex_out = 0;
  fading_dir = strdup(DEFAULT_FADING_DIR);
  user_multipath_delay = -1.0;

  int o = 0;
  int opt_idx = 0;
  while (o != -1) {
    static struct option long_opts[] = {
        {"complexout", no_argument, 0, 'o'},
        {"ctest", no_argument, 0, 't'},
        {"clip", required_argument, 0, 'c'},
        {"fading_dir", required_argument, 0, 'u'},
        {"freq", required_argument, 0, 'f'},
        {"Fs", required_argument, 0, 'r'},
        {"gain", required_argument, 0, 'g'},
        {"ssbfilt", required_argument, 0, 's'},
        {"help", no_argument, 0, 'h'},
        {"mpg", no_argument, 0, 'i'},
        {"mpp", no_argument, 0, 'p'},
        {"mpd", no_argument, 0, 'd'},
        {"multipath_delay", required_argument, 0, 'm'},
        {"No", required_argument, 0, 'n'},
        {0, 0, 0, 0}};

    o = getopt_long(argc, argv, "c:df:g:im:n:opr:s:tu:h", long_opts, &opt_idx);

    switch (o) {
      case 'c':
        clip = atof(optarg);
        break;
      case 'd':
        fading_en = 3;
        break;
      case 'f':
        foff_hz = atof(optarg);
        break;
      case 'g':
        gain = atof(optarg);
        break;
      case 'i':
        fading_en = 1;
        break;
      case 'm':
        user_multipath_delay = atof(optarg);
        break;
      case 'n':
        NodB = atof(optarg);
        break;
      case 'o':
        complex_out = 1;
        break;
      case 'p':
        fading_en = 2;
        break;
      case 'r':
        Fs = atoi(optarg);
        break;
      case 's':
        ssbfilt_en = atoi(optarg);
        break;
      case 't':
        ctest = 1;
        break;
      case 'u':
        fading_dir = strdup(optarg);
        break;
      case 'h':
      case '?':
        goto helpmsg;
        break;
    }
  }

  phase_ch.real = 1.0;
  phase_ch.imag = 0.0;

  /*  N = var = NoFs */

  // arbitrary noise scaling, to maintain backwards compatibility with many
  // tests.  TODO make the No units more sensible, and fix all the tests that
  // depend on this scaling
  No = pow(10.0, NodB / 10.0) * 1000 * 1000;
  variance = Fs * No;

  tx_pwr = tx_pwr_fade = noise_pwr = 0.0;
  noutclipped = 0;
  nclipped = 0;
  peak = 0.0;

  /* init HF fading model */

  ffading = NULL;
  nhfdelay = 0;
  if (fading_en) {
    char fname[256];

    if (fading_en == 1) {
      sprintf(fname, "%s/%s", fading_dir, MPG_FADING_FILE_NAME);
      ffading = fopen(fname, "rb");
      if (ffading == NULL) {
      cant_load_fading_file:
        fprintf(stderr,
                "-----------------------------------------------------\n");
        fprintf(stderr, "ch ERROR: Can't find fading file: %s\n", fname);
        fprintf(
            stderr,
            "\nAdjust path --fading_dir or use GNU Octave to generate:\n\n");
      gen_fading_file:
        fprintf(stderr, "$ octave-cli\n");
        fprintf(stderr, "octave:24> pkg load signal\n");
        fprintf(stderr, "octave:24> time_secs=60\n");
        fprintf(stderr,
                "octave:25> ch_fading(\"faster_fading_samples.float\", 8000, "
                "2.0, 8000*time_secs)\n");
        fprintf(stderr,
                "octave:26> ch_fading(\"fast_fading_samples.float\", 8000, "
                "1.0, 8000*time_secs)\n");
        fprintf(stderr,
                "octave:27> ch_fading(\"slow_fading_samples.float\", 8000, "
                "0.1, 8000*time_secs)\n");
        fprintf(stderr,
                "-----------------------------------------------------\n");
        exit(1);
      }
      nhfdelay = floor(MPG_DELAY_MS * Fs / 1000);
    }

    if (fading_en == 2) {
      sprintf(fname, "%s/%s", fading_dir, MPP_FADING_FILE_NAME);
      ffading = fopen(fname, "rb");
      if (ffading == NULL) goto cant_load_fading_file;
      nhfdelay = floor(MPP_DELAY_MS * Fs / 1000);
    }

    if (fading_en == 3) {
      sprintf(fname, "%s/%s", fading_dir, MPD_FADING_FILE_NAME);
      ffading = fopen(fname, "rb");
      if (ffading == NULL) goto cant_load_fading_file;
      nhfdelay = floor(MPD_DELAY_MS * Fs / 1000);
    }

    ch_fdm_delay = (COMP *)MALLOC((nhfdelay + COHPSK_NOM_SAMPLES_PER_FRAME) *
                                  sizeof(COMP));
    assert(ch_fdm_delay != NULL);
    for (i = 0; i < nhfdelay + COHPSK_NOM_SAMPLES_PER_FRAME; i++) {
      ch_fdm_delay[i].real = 0.0;
      ch_fdm_delay[i].imag = 0.0;
    }

    /* optionally override delay from command line */
    if (user_multipath_delay >= 0.0)
      nhfdelay = floor(user_multipath_delay * Fs / 1000);

    /* first values in file are HF gains */

    for (i = 0; i < 4; i++) ret = fread(&hf_gain, sizeof(float), 1, ffading);
    // fprintf(stderr, "hf_gain: %f\n", hf_gain);
  }

  assert(HT_N == sizeof(ht_coeff) / sizeof(COMP));
  for (i = 0; i < HT_N; i++) {
    htbuf[i] = 0.0;
  }
  for (i = 0; i < SSBFILT_N; i++) {
    ssbfiltbuf[i].real = 0.0;
    ssbfiltbuf[i].imag = 0.0;
  }
  COMP lo_phase = {1.0, 0.0};
  COMP lo_freq;
  lo_freq.real = cos(2.0 * M_PI * SSBFILT_CENTRE / Fs);
  lo_freq.imag = sin(2.0 * M_PI * SSBFILT_CENTRE / Fs);

  fprintf(stderr,
          "ch: Fs: %d NodB: %4.2f foff: %4.2f Hz fading: %d nhfdelay: %d clip: "
          "%4.2f ssbfilt: %d complexout: %d\n",
          Fs, NodB, foff_hz, fading_en, nhfdelay, clip, ssbfilt_en,
          complex_out);

  /* --------------------------------------------------------*\
                                Main Loop
  \*---------------------------------------------------------*/

  frames = 0;
  while (fread(buf, sizeof(short), BUF_N, fin) == BUF_N) {
    frames++;

    /* Hilbert Transform to produce complex signal so we can do
       single sided freq shifts, HF channel models, and analog
       compression.  Allows us to use real signal I/O.

       As the real and imag filters both have unity gain, ch_in[]
       has twice the power of the real input signal buf[].
    */

    for (i = 0, j = HT_N; i < BUF_N; i++, j++) {
      htbuf[j] = (float)buf[i] * gain;

      /* FIR filter with HT to get imag, just delay to get real */

      ch_in[i].real = 0.0;
      ch_in[i].imag = 0.0;
      for (k = 0; k < HT_N; k++) {
        ch_in[i].real += htbuf[j - k] * ht_coeff[k].real;
        ch_in[i].imag += htbuf[j - k] * ht_coeff[k].imag;
      }
      // printf("%d %f %f\n", i, ch_in[i].real, ch_in[i].imag);
    }
    assert(j <= (BUF_N + HT_N));

    /* update HT memory */
    for (i = 0; i < HT_N; i++) htbuf[i] = htbuf[i + BUF_N];

    /* --------------------------------------------------------*\
                        Clipping  mag of complex signal
    \*---------------------------------------------------------*/

    for (i = 0; i < BUF_N; i++) {
      float mag =
          sqrt(ch_in[i].real * ch_in[i].real + ch_in[i].imag * ch_in[i].imag);
      float angle = atan2(ch_in[i].imag, ch_in[i].real);
      if (mag > clip) {
        mag = clip;
        nclipped++;
      }
      tx_pwr += mag * mag;
      /* we get a bit of overshoot in peak measurements if HT filter hasn't been
       * primed */
      if (frames * BUF_N > HT_N)
        if (mag > peak) {
          peak = mag;
          // fprintf(stderr, "%d %f\n",frames, mag);
        }
      ch_in[i].real = mag * cos(angle);
      ch_in[i].imag = mag * sin(angle);
    }

    /* --------------------------------------------------------*\
                             Channel
    \*---------------------------------------------------------*/

    fdmdv_freq_shift_coh(ch_fdm, ch_in, foff_hz, Fs, &phase_ch, BUF_N);

    /* optional HF fading -------------------------------------*/

    if (fading_en) {
      /* update delayed signal buffer */

      for (i = 0; i < nhfdelay; i++) ch_fdm_delay[i] = ch_fdm_delay[i + BUF_N];
      for (j = 0; j < BUF_N; i++, j++) ch_fdm_delay[i] = ch_fdm[j];

      /* combine direct and delayed paths, both multiplied by
         "spreading" (Doppler) functions */

      for (i = 0; i < BUF_N; i++) {
        ret = fread(&aspread, sizeof(COMP), 1, ffading);
        if (ret == 0) {
          fprintf(stderr,
                  "ch: Fading file finished - simulation stopping.  You may "
                  "need more samples:\n");
          goto gen_fading_file;
        }
        ret = fread(&aspread_2ms, sizeof(COMP), 1, ffading);
        if (ret == 0) {
          fprintf(stderr,
                  "ch: Fading file finished - simulation stopping.  You may "
                  "need more samples:\n");
          goto gen_fading_file;
        }
        // printf("%f %f %f %f\n", aspread.real, aspread.imag, aspread_2ms.real,
        // aspread_2ms.imag);

        direct = cmult(aspread, ch_fdm[i]);
        delayed = cmult(aspread_2ms, ch_fdm_delay[i]);
        ch_fdm[i] = fcmult(hf_gain, cadd(direct, delayed));
      }
    }

    /* Measure power after fading model to make sure average pwr
       is the same as AWGN channels. We only output the real
       signal, which is half the power. */

    for (i = 0; i < BUF_N; i++) {
      tx_pwr_fade += pow(ch_fdm[i].real, 2.0);
    }

    /* AWGN noise ------------------------------------------*/

    for (i = 0; i < BUF_N; i++) {
      COMP n = noise();
      scaled_noise = fcmult(sqrt(variance), n);
      ch_fdm[i] = cadd(ch_fdm[i], scaled_noise);
      noise_pwr += pow(scaled_noise.real, 2.0) + pow(scaled_noise.imag, 2.0);
    }

    /* FIR filter to simulate (a rather flat) SSB filter. We
       filter the complex signal by shifting it down to DC and
       using real coefficients. */

    for (i = 0, j = SSBFILT_N; i < BUF_N; i++, j++) {
      if (ssbfilt_en) {
        ssbfiltbuf[j] = cmult(ch_fdm[i], cconj(lo_phase));
        ssbfiltout[i].real = 0.0;
        ssbfiltout[i].imag = 0.0;
        for (k = 0; k < SSBFILT_N; k++) {
          ssbfiltout[i].real += ssbfiltbuf[j - k].real * ssbfilt_coeff[k];
          ssbfiltout[i].imag += ssbfiltbuf[j - k].imag * ssbfilt_coeff[k];
        }
        ssbfiltout[i] = cmult(ssbfiltout[i], lo_phase);
        lo_phase = cmult(lo_phase, lo_freq);
      } else {
        ssbfiltout[i] = ch_fdm[i];
      }
    }

    /* update SSB filter memory */
    for (i = 0; i < SSBFILT_N; i++) ssbfiltbuf[i] = ssbfiltbuf[i + BUF_N];

    int nout = (complex_out + 1) * BUF_N;
    short bufout[nout], *pout = bufout;
    for (i = 0; i < BUF_N; i++) {
      sam = ssbfiltout[i].real;
      if (sam > 32767.0) {
        noutclipped++;
        sam = 32767.0;
      }
      if (sam < -32767.0) {
        noutclipped++;
        sam = -32767.0;
      }
      *pout++ = sam;
      if (complex_out) {
        sam = ssbfiltout[i].imag;
        if (sam > 32767.0) {
          noutclipped++;
          sam = 32767.0;
        }
        if (sam < -32767.0) {
          noutclipped++;
          sam = -32767.0;
        }
        *pout++ = sam;
      }
    }

    fwrite(bufout, sizeof(short), nout, fout);

    /* if this is in a pipeline, we probably don't want the usual
       buffering to occur */

    if (fout == stdout) fflush(stdout);
  }

  fclose(fin);
  fclose(fout);

  int nsamples = frames * BUF_N;
  papr = 10 * log10(peak * peak / (tx_pwr / nsamples));
  CNo = 10 * log10(tx_pwr / (noise_pwr / (Fs)));
  snr3k = CNo - 10 * log10(3000);
  float outclipped_percent = noutclipped * 100.0 / nsamples;
  fprintf(stderr, "ch: SNR3k(dB): %8.2f  C/No....: %8.2f\n", snr3k, CNo);
  fprintf(stderr,
          "ch: peak.....: %8.2f  RMS.....: %8.2f   CPAPR.....: %5.2f \n", peak,
          sqrt(tx_pwr / nsamples), papr);
  fprintf(stderr,
          "ch: Nsamples.: %8d  clipped.: %8.2f%%  OutClipped: %5.2f%%\n",
          nsamples, nclipped * 100.0 / nsamples, outclipped_percent);
  if (outclipped_percent > 0.1)
    fprintf(stderr, "ch: WARNING output clipping\n");

  if (ffading != NULL) fclose(ffading);
  if (ch_fdm_delay != NULL) FREE(ch_fdm_delay);
  if (ctest) {
    /* special ctest mode: check CPAPR is around 0dB */
    if (fabs(papr) < 0.7)
      return 0;
    else
      return 1;
  } else
    return 0;
}