/**
 *
 * \section COPYRIGHT
 *
 * Copyright 2013-2014 The libLTE Developers. See the
 * COPYRIGHT file at the top-level directory of this distribution.
 *
 * \section LICENSE
 *
 * This file is part of the libLTE library.
 *
 * libLTE is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation, either version 3 of
 * the License, or (at your option) any later version.
 *
 * libLTE 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 Lesser General Public License for more details.
 *
 * A copy of the GNU Lesser General Public License can be found in
 * the LICENSE file in the top-level directory of this distribution
 * and at http://www.gnu.org/licenses/.
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <unistd.h>
#include <math.h>
#include <sys/time.h>
#include <unistd.h>
#include <assert.h>
#include <signal.h>

#include "liblte/phy/phy.h"

#ifndef DISABLE_UHD
#include "liblte/cuhd/cuhd.h"
void *uhd;
#endif

#ifndef DISABLE_GRAPHICS
#include "liblte/graphics/plot.h"
plot_real_t poutfft;
plot_complex_t pce;
plot_scatter_t pscatrecv, pscatequal;
#endif

#define MHZ       1000000
#define SAMP_FREQ     1920000
#define FLEN       9600
#define FLEN_PERIOD    0.005

#define NOF_PORTS 2

float find_threshold = 20.0;
int max_track_lost = 20, nof_frames = -1;
int track_len = 300;
char *input_file_name = NULL;
int disable_plots = 0;

int go_exit = 0;

float uhd_freq = 2600000000.0, uhd_gain = 20.0;
char *uhd_args = "";

filesource_t fsrc;
cf_t *input_buffer, *fft_buffer, *ce[MAX_PORTS];
pbch_t pbch;
lte_fft_t fft;
chest_t chest;
sync_t ssync;
cfo_t cfocorr;

enum sync_state {
  FIND, TRACK
};

void usage(char *prog) {
  printf("Usage: %s [iagfndvp]\n", prog);
  printf("\t-i input_file [Default use USRP]\n");
#ifndef DISABLE_UHD
  printf("\t-a UHD args [Default %s]\n", uhd_args);
  printf("\t-g UHD RX gain [Default %.2f dB]\n", uhd_gain);
  printf("\t-f UHD RX frequency [Default %.1f MHz]\n", uhd_freq / 1000000);
#else
  printf("\t   UHD is disabled. CUHD library not available\n");
#endif
  printf("\t-n nof_frames [Default %d]\n", nof_frames);
  printf("\t-p PSS threshold [Default %f]\n", find_threshold);
#ifndef DISABLE_GRAPHICS
  printf("\t-d disable plots [Default enabled]\n");
#else
  printf("\t plots are disabled. Graphics library not available\n");
#endif
  printf("\t-v [set verbose to debug, default none]\n");
}

void parse_args(int argc, char **argv) {
  int opt;
  while ((opt = getopt(argc, argv, "iagfndvp")) != -1) {
    switch (opt) {
    case 'i':
      input_file_name = argv[optind];
      break;
    case 'a':
      uhd_args = argv[optind];
      break;
    case 'g':
      uhd_gain = atof(argv[optind]);
      break;
    case 'f':
      uhd_freq = atof(argv[optind]);
      break;
    case 'p':
      find_threshold = atof(argv[optind]);
      break;
    case 'n':
      nof_frames = atoi(argv[optind]);
      break;
    case 'd':
      disable_plots = 1;
      break;
    case 'v':
      verbose++;
      break;
    default:
      usage(argv[0]);
      exit(-1);
    }
  }
}

#ifndef DISABLE_GRAPHICS

void init_plots() {
  plot_init();
  plot_real_init(&poutfft);
  plot_real_setTitle(&poutfft, "Output FFT - Magnitude");
  plot_real_setLabels(&poutfft, "Index", "dB");
  plot_real_setYAxisScale(&poutfft, -60, 0);
  plot_real_setXAxisScale(&poutfft, 1, 504);

  plot_complex_init(&pce);
  plot_complex_setTitle(&pce, "Channel Estimates");
  plot_complex_setYAxisScale(&pce, Ip, -0.01, 0.01);
  plot_complex_setYAxisScale(&pce, Q, -0.01, 0.01);
  plot_complex_setYAxisScale(&pce, Magnitude, 0, 0.01);
  plot_complex_setYAxisScale(&pce, Phase, -M_PI, M_PI);

  plot_scatter_init(&pscatrecv);
  plot_scatter_setTitle(&pscatrecv, "Received Symbols");
  plot_scatter_setXAxisScale(&pscatrecv, -0.01, 0.01);
  plot_scatter_setYAxisScale(&pscatrecv, -0.01, 0.01);

  plot_scatter_init(&pscatequal);
  plot_scatter_setTitle(&pscatequal, "Equalized Symbols");
  plot_scatter_setXAxisScale(&pscatequal, -1, 1);
  plot_scatter_setYAxisScale(&pscatequal, -1, 1);
}

#endif

int base_init(int frame_length) {
  int i;

#ifndef DISABLE_GRAPHICS
  if (!disable_plots) {
    init_plots();
  }
#else
  printf("-- PLOTS are disabled. Graphics library not available --\n\n");
#endif

  if (input_file_name) {
    if (filesource_init(&fsrc, input_file_name, COMPLEX_FLOAT_BIN)) {
      return -1;
    }
  } else {
    /* open UHD device */
#ifndef DISABLE_UHD
    printf("Opening UHD device...\n");
    if (cuhd_open(uhd_args, &uhd)) {
      fprintf(stderr, "Error opening uhd\n");
      return -1;
    }
#else
    printf("Error UHD not available. Select an input file\n");
    return -1;
#endif
  }

  input_buffer = (cf_t*) malloc(frame_length * sizeof(cf_t));
  if (!input_buffer) {
    perror("malloc");
    return -1;
  }

  fft_buffer = (cf_t*) malloc(CPNORM_NSYMB * 72 * sizeof(cf_t));
  if (!fft_buffer) {
    perror("malloc");
    return -1;
  }
  for (i = 0; i < MAX_PORTS; i++) {
    ce[i] = (cf_t*) malloc(CPNORM_NSYMB * 72 * sizeof(cf_t));
    if (!ce[i]) {
      perror("malloc");
      return -1;
    }
  }
  if (sync_init(&ssync, FLEN)) {
    fprintf(stderr, "Error initiating PSS/SSS\n");
    return -1;
  }
  if (chest_init(&chest, LINEAR, CPNORM, 6, NOF_PORTS)) {
    fprintf(stderr, "Error initializing equalizer\n");
    return -1;
  }

  if (cfo_init(&cfocorr, FLEN)) {
    fprintf(stderr, "Error initiating CFO\n");
    return -1;
  }
  if (lte_fft_init(&fft, CPNORM, 6)) {
    fprintf(stderr, "Error initializing FFT\n");
    return -1;
  }

  return 0;
}

void base_free() {
  int i;

  if (input_file_name) {
    filesource_free(&fsrc);
  } else {
#ifndef DISABLE_UHD
    cuhd_close(uhd);
#endif
  }

#ifndef DISABLE_GRAPHICS
  plot_exit();
#endif

  sync_free(&ssync);
  lte_fft_free(&fft);
  chest_free(&chest);
  cfo_free(&cfocorr);

  free(input_buffer);
  free(fft_buffer);
  for (i = 0; i < MAX_PORTS; i++) {
    free(ce[i]);
  }
}

int mib_decoder_init(int cell_id) {

  lte_cell_t cell; 
  
  cell.id = cell_id;
  cell.nof_ports = MAX_PORTS;
  cell.nof_prb = 6;
  cell.cp = CPNORM;
  
  if (chest_ref_LTEDL(&chest, cell)) {
    fprintf(stderr, "Error initializing reference signal\n");
    return -1;
  }

  if (pbch_init(&pbch, 6, cell_id, CPNORM)) {
    fprintf(stderr, "Error initiating PBCH\n");
    return -1;
  }
  DEBUG("PBCH initiated cell_id=%d\n", cell_id);
  return 0;
}

int mib_decoder_run(cf_t *input, pbch_mib_t *mib) {
  int i, n;
  lte_fft_run_slot(&fft, input, fft_buffer);

  /* Get channel estimates for each port */
  for (i = 0; i < NOF_PORTS; i++) {
    chest_ce_slot_port(&chest, fft_buffer, ce[i], 1, i);
  }

  DEBUG("Decoding PBCH\n", 0);
  n = pbch_decode(&pbch, fft_buffer, ce, mib);

#ifndef DISABLE_GRAPHICS
  float tmp[72 * 7];
  if (!disable_plots) {
    for (i = 0; i < 72 * 7; i++) {
      tmp[i] = 10 * log10f(cabsf(fft_buffer[i]));
      if (isinf(tmp[i])) {
        tmp[i] = -80;
      }
    }

    plot_real_setNewData(&poutfft, tmp, 72 * 7);
    plot_complex_setNewData(&pce, ce[0], 72 * 7);
    plot_scatter_setNewData(&pscatrecv, pbch.pbch_symbols[0], pbch.nof_symbols);
    if (n) {
      plot_scatter_setNewData(&pscatequal, pbch.pbch_d, pbch.nof_symbols);
    }
  }
#endif

  return n;
}

void sigintHandler(int sig_num) {
  go_exit = 1;
}

int main(int argc, char **argv) {
  int frame_cnt;
  int cell_id;
  int find_idx, track_idx, last_found;
  enum sync_state state;
  int nslot;
  pbch_mib_t mib;
  float cfo;
  int n;
  int nof_found_mib = 0;
  float timeoffset = 0;

#ifdef DISABLE_UHD
  if (argc < 3) {
    usage(argv[0]);
    exit(-1);
  }
#endif

  parse_args(argc, argv);

  if (base_init(FLEN)) {
    fprintf(stderr, "Error initializing memory\n");
    exit(-1);
  }

  sync_pss_det_peak_to_avg(&ssync);

  if (!input_file_name) {
#ifndef DISABLE_UHD
    INFO("Setting sampling frequency %.2f MHz\n", (float) SAMP_FREQ/MHZ);
    cuhd_set_rx_srate(uhd, SAMP_FREQ);
    cuhd_set_rx_gain(uhd, uhd_gain);
    /* set uhd_freq */
    cuhd_set_rx_freq(uhd, (double) uhd_freq);
    cuhd_rx_wait_lo_locked(uhd);
    DEBUG("Set uhd_freq to %.3f MHz\n", (double ) uhd_freq);

    DEBUG("Starting receiver...\n", 0);
    cuhd_start_rx_stream(uhd);
#endif
  }

  printf("\n --- Press Ctrl+C to exit --- \n");
  signal(SIGINT, sigintHandler);

  state = FIND;
  nslot = 0;
  find_idx = 0;
  cfo = 0;
  mib.sfn = -1;
  frame_cnt = 0;
  last_found = 0;
  sync_set_threshold(&ssync, find_threshold);

  while (!go_exit && (frame_cnt < nof_frames || nof_frames == -1)) {
    INFO(" -----   RECEIVING %d SAMPLES ---- \n", FLEN);
    if (input_file_name) {
      n = filesource_read(&fsrc, input_buffer, FLEN);
      if (n == -1) {
        fprintf(stderr, "Error reading file\n");
        exit(-1);
      } else if (n < FLEN) {
        filesource_seek(&fsrc, 0);
        filesource_read(&fsrc, input_buffer, FLEN);
      }
    } else {
#ifndef DISABLE_UHD
      cuhd_recv(uhd, input_buffer, FLEN, 1);
#endif
    }

    switch (state) {
    case FIND:
      /* find peak in all frame */
      find_idx = sync_find(&ssync, input_buffer);
      INFO("FIND %3d:\tPAR=%.2f\n", frame_cnt, sync_get_peak_to_avg(&ssync));
      if (find_idx != -1) {
        /* if found peak, go to track and set track threshold */
        cell_id = sync_get_cell_id(&ssync);
        if (cell_id != -1) {
          frame_cnt = -1;
          last_found = 0;
          mib_decoder_init(cell_id);
          nof_found_mib = 0;
          nslot = sync_get_slot_id(&ssync);
          nslot = (nslot + 10) % 20;
          cfo = 0;
          timeoffset = 0;
          printf("\n");
          state = TRACK;
        } else {
          printf("cellid=-1\n");
        }
      }
      if (verbose == VERBOSE_NONE) {
        printf("Finding PSS... PAR=%.2f\r", sync_get_peak_to_avg(&ssync));
      }
      break;
    case TRACK:
      /* Find peak around known position find_idx */
      INFO("TRACK %3d: PSS find_idx %d offset %d\n", frame_cnt, find_idx,
          find_idx - track_len);
      track_idx = sync_track(&ssync, &input_buffer[find_idx - track_len]);

      if (track_idx != -1) {
        /* compute cumulative moving average CFO */
        cfo = (sync_get_cfo(&ssync) + frame_cnt * cfo) / (frame_cnt + 1);
        /* compute cumulative moving average time offset */
        timeoffset = (float) (track_idx - track_len + timeoffset * frame_cnt)
            / (frame_cnt + 1);
        last_found = frame_cnt;
        find_idx = (find_idx + track_idx - track_len) % FLEN;
      } else {
        /* if sync not found, adjust time offset with the averaged value */
        find_idx = (find_idx + (int) timeoffset) % FLEN;
      }

      /* if we missed too many PSS go back to FIND */
      if (frame_cnt - last_found > max_track_lost) {
        INFO("%d frames lost. Going back to FIND", frame_cnt - last_found);
        printf("\r\n");
        fflush(stdout);
        printf("\r\n");
        state = FIND;
      }

      // Correct CFO
      INFO("Correcting CFO=%.4f\n", cfo);

      cfo_correct(&cfocorr, input_buffer, -cfo / 128);

      if (nslot == 0 && find_idx + 960 < FLEN) {
        INFO("Finding MIB at idx %d\n", find_idx);
        if (mib_decoder_run(&input_buffer[find_idx], &mib)) {
          INFO("MIB detected attempt=%d\n", frame_cnt);
          if (verbose == VERBOSE_NONE) {
            if (!nof_found_mib) {
              printf("\r\n");
              fflush(stdout);
              printf("\r\n");
              printf(" - Phy. CellId:\t%d\n", cell_id);
              pbch_mib_fprint(stdout, &mib);
            }
          }
          nof_found_mib++;
        } else {
          INFO("MIB not found attempt %d\n", frame_cnt);
        }
        if (frame_cnt) {
          printf(
              "SFN: %4d, CFO: %+.4f KHz, SFO: %+.4f Khz, TimeOffset: %4d, Errors: %4d/%4d, ErrorRate: %.1e\r",
              mib.sfn, cfo * 15, timeoffset / 5, find_idx,
              frame_cnt - 2 * (nof_found_mib - 1), frame_cnt,
              (float) (frame_cnt - 2 * (nof_found_mib - 1)) / frame_cnt);
          fflush(stdout);
        }
      }
      if (input_file_name) {
        usleep(5000);
      }
      nslot = (nslot + 10) % 20;
      break;
    }
    frame_cnt++;
  }

  base_free();

  printf("\nBye\n");
  exit(0);
}