// Copyright 2003 Regents of the University of California

// Astropulse is free software; you can redistribute it and/or modify it under
// the terms of the GNU General Public License as published by the Free
// Software Foundation; either version 2, or (at your option) any later
// version.

// Astropulse 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 General Public License along
// with Astropulse; see the file COPYING.  If not, write to the Free Software
// Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

/* ap_science.C */
/* Functions for AstroPulse scientific computations */

#ifndef _WIN32
#include "ap_config.h"
#else
#include "boinc_win.h"
#endif

#include <cmath>
#include <cassert>
#include "astropulse.h"
#include "boinc_api.h"
#include "sbtf.h"
#include "ap_debug.h"
#include "fftw3.h"

/* Windows is stupid */
#ifndef M_PI
#define M_PI 3.141592653589793238462
#endif

namespace Astropulse {
const char* WISDOM_FNAME =  "wisdom.dat";

inline void Science::Wisdom::load() {
  FILE *wisdom;
  if (wisdom=boinc_fopen(WISDOM_FNAME,"r")) {
    char *wiz=(char *)calloc(1024,64);
    int n=0;
    while (wiz && n<64*1024 && !feof(wisdom)) {
      n+=fread(wiz+n,1,80,wisdom);
    }
    fftwf_import_wisdom_from_string(wiz);
    free(wiz);
    fclose(wisdom);
  }
}

inline void Science::Wisdom::save() {
  // Save new wisdom
  FILE *wisdom=boinc_fopen(WISDOM_FNAME,"w");
  if (wisdom) {
    char *wiz=fftwf_export_wisdom_to_string();
    if (wiz) {
      fwrite(wiz,strlen(wiz),1,wisdom);
    }
    fclose(wisdom);
  }
}

// Allocate memory, and create fft plans
void Science::init() {
  const size_t& fft_len = client.ap_shmem->ap_gdata.state.fft_len;
  int fft_len2 = 1024;
  fftlen_envelope = 128;
  fftlen_patch = 512;

  data  = (fftwf_complex *)fftwf_malloc(sizeof(fftwf_complex)*fft_len);
  temp  = (fftwf_complex *)fftwf_malloc(sizeof(fftwf_complex)*fft_len);
  chirp  = (fftwf_complex *)fftwf_malloc(sizeof(fftwf_complex)*fft_len*client.ap_shmem->ap_gdata.state.dm_chunk_large);
  patch = (fftwf_complex *)fftwf_malloc(sizeof(fftwf_complex)*fftlen_patch);
  envelope_part = (fftwf_complex *) fftwf_malloc(sizeof(fftwf_complex)*fftlen_envelope);

  
  envelope = new double[fftlen_envelope];
  pre_envelope = new double[fftlen_envelope];

  client.status = "Creating FFT plans.";

  wisdom.load();

  client.status = "Creating forward FFT plan.";

  fp = fftwf_plan_dft_1d(fft_len, data, data, FFTW_FORWARD, FFTW_MEASURE);
  fp2 = fftwf_plan_dft_1d(fft_len2, data, data, FFTW_FORWARD, FFTW_MEASURE);
  fp_patch_inv = fftwf_plan_dft_1d(fftlen_patch, patch, patch, FFTW_BACKWARD, FFTW_MEASURE);
  fp_envelope = fftwf_plan_dft_1d(fftlen_envelope, envelope_part, envelope_part, FFTW_FORWARD, FFTW_MEASURE);


  client.status = "Creating inverse FFT plan.";
  ip = fftwf_plan_dft_1d(fft_len, temp, temp, FFTW_BACKWARD, FFTW_MEASURE);

  wisdom.save();

}

void Science::finish() {
  fftwf_destroy_plan(fp);
  fftwf_destroy_plan(ip);
  fftwf_destroy_plan(fp_patch_inv);

  fftwf_free(chirp);
  fftwf_free(temp);
  fftwf_free(data);
  fftwf_free(patch);
}

/* Science::build_chirp_table
 *
 * arguments:
 * int dm_start: the first dm to use in the chirp table.
 *  The chirp table will use all dms from dm_start to
 * dm_start + client.ap_shmem->ap_gdata.state.dm_chunk - 1. 
 *
 * This function creates an array of complex numbers 
 * (of type fftw_complex) which
 * represent a chirp.  The array indices represent frequencies, so
 * the chirp is described in frequency space.  We will multiply the
 * chirp by the frequency-space representation of the incoming signal
 * in order to dechirp the signal, so it is really a "dechirp" table.
 * That is, we use the reciprocal of the chirp, instead of the chirp
 * itself.  Since the chirp's amplitude is 1 everywhere, taking its
 * reciprocal is a simple matter of taking its complex conjugate.
 * 
 * The first <state.fft_len> complex numbers correspond to the
 * dm indicated by dm_start.  The next set of numbers correspond to
 * dm_start+1, and so on until dm_start + client.ap_shmem->ap_gdata.state.dm_chunk - 1.
 */
void Science::build_chirp_table(int dm_start) {
  int i,j;
  double freq;
  double phase;

  client.status = "Building chirp table.";

  state_t& state = client.ap_shmem->ap_gdata.state;

  /* For debugging, see below
  static int idum = -1;
  */

  // phase function determined by method of steepest descent.
  // 0.00176 = 2.5 MHz / 1.42 GHz.  When multiplied by the maximum
  // value freq = 0.5, it will give 1.25 MHz / 1.42 GHz.
  for (i=0; i<state.dm_chunk_large; i++) {
    for (j=0; j<state.fft_len; j++) {
      // To get exact chirp function, we've multiplied the phase
      // by 1 / (1 + 0.00176 * freq) from the linear chirp function
      // This is equal to nu0 / nu, where nu0 = 1.42 GHz and
      // nu = 1.42 GHz + freq in Hz = 1.42 GHz + (freq / (4 * 10^-7)) Hz,
      // resulting in the value 0.00176 = 1 / (1.42 * 10^9 * 4 * 10^-7)
      freq = (double)j/(double)state.fft_len;
      freq = (j<state.fft_len/2)?freq:freq-1.0;
      phase = M_PI*(double)(dm_start+i)*freq*freq*(1 / (1 + 0.00176 * freq));

      /* For debugging: use a random phase instead of a chirp
      phase = 2 * M_PI * ran1(idum);
      */

      chirp[i*state.fft_len + j][0] = (float)cos(phase);
      chirp[i*state.fft_len + j][1] = (float)sin(phase);
    }
  }

}

void Science::convert_bits_to_float() {
  client.status = "Converting raw data.";
  if (splitter_bits_to_float(&client.rawdata[client.ap_shmem->ap_gdata.state.data_chunk_now],
                             (float *)data, client.ap_shmem->ap_gdata.state.fft_len)) {
    exit(SBTF_ERR);
  }
}

void Science::compute_forward_fft() {
  client.status = "Computing forward FFT.";
  fftwf_execute(fp);
}

void Science::high_pass_filter() {
  /* One frequency bin is 2.5 MHz / fft_len worth of frequency.
   * So the width of the Gaussian filter in bins is:
   * (Note the sample rate is stored in the recorder_config, which is not part of the
   *  client.ap_shmem->ap_gdata.state.)
   */
  const size_t& fft_len = client.ap_shmem->ap_gdata.state.fft_len;
  if (client.ap_shmem->ap_gdata.state.high_pass > 0) {
    float sigma = client.ap_shmem->ap_gdata.state.high_pass * ((float) fft_len / 2500000);
    int i;
    int nf;  // For negative frequencies
    float Gaussian, multiplier;
    static float normalize = 1 / sqrt(1 - sqrt(2 * 3.1415926) * sigma / fft_len);

    for (i = 0; i < (int) fft_len; i++) {
      if (i < fft_len/2) {
        Gaussian = exp(-i * i / (2 * sigma * sigma));
      } else {
        nf = fft_len - i;
        Gaussian = exp(-nf * nf / (2 * sigma * sigma));
      }
      /* Multiply the power by 1 - Gaussian. */
      multiplier = sqrt(1 - Gaussian);
      /* Conserve total noise power. */
      multiplier *= normalize;
      data[i][0] *= multiplier;
      data[i][1] *= multiplier;
    }
  }
}

/* Science::dechirp
 * Now DIVIDES by the chirp function, instead of multiplying.
 *
 * Arguments:
 * dm: dispersion measure to dechirp by, in units of 0.4 microseconds
 * power: array of length fft_len/2, to hold the power in each time bin
 *        I'm not sure why we only record the first fft_len/2 bins.
 * sign: +1 for a positive dm, -1 for a negative dm
 */
void Science::dechirp(int dm, float *power, int sign) {
  const size_t& fft_len = client.ap_shmem->ap_gdata.state.fft_len;
  int n = dm*fft_len;
  size_t i;
  static float avg_pow = 2*fft_len*float(fft_len);
  // avg_pow contains N^2 for unnormalized fft, and 2 for 1-bit
  // convention

  client.status = "Dechirping.";
  /* For debugging.  Actually I don't think this runs fast enough
  static int idum = -1;
  for(i=0; i< fft_len; i++) {
   data[i][0] = random_discrete_sum(fft_len, idum);
   data[i][1] = random_discrete_sum(fft_len, idum);
  } */

  // Multiply by the complex conjugate of chirp, so we are 
  // DIVIDING by the chirp function.
  for (i=0; i<fft_len; i++) {
    /* positive DM */
    if (sign >= 0) {
      temp[i][0] = data[i][0]*chirp[n+i][0] + data[i][1]*chirp[n+i][1];
      temp[i][1] = data[i][1]*chirp[n+i][0] - data[i][0]*chirp[n+i][1];
    }
    /* negative DM */
    else {
      temp[i][0] = data[i][0]*chirp[n+i][0] - data[i][1]*chirp[n+i][1];
      temp[i][1] = data[i][1]*chirp[n+i][0] + data[i][0]*chirp[n+i][1];
    }
  }

  client.status = "Computing inverse FFT.";
  fftwf_execute(ip);

  /* only save 1/2 of the fft ... need to check this */
  for (i=0; i<fft_len/2; i++) {
    client.power[i] = (float)((temp[i][0]*temp[i][0] + temp[i][1]*temp[i][1])/avg_pow);
  }
}

/* coadd_in_place
 *
 * Arguments:
 * power: an array of power values, obtained after dechirping
 * size: the number of floats in the initial power array.
 * (After the coadd, the number of floats will be halved.)
 *
 * This function halves the size of the power array by adding
 * array elements in consecutive pairs.
 */
void Science::coadd_in_place(float *power, int size) {
  int i;

  for (i=0; i<size/2; i++) {
    power[i] = power[2*i] + power[2*i+1];
  }
}
}

