#include "astropulse.h"
#include "ap_graphics.h"
#include "fftw3.h"
#include "sbtf.h"
#include "ap_debug.h"
#include <stdio.h>
#include "mtrand.h"
#include <complex>

using std::complex;
// #include "seti_dr2filter.h"

namespace Astropulse {

void Science::print_DC_strengths() {

  const int NUM_STRENGTHS = 8;
  float strengths[NUM_STRENGTHS];
  float totals[NUM_STRENGTHS];
  float averages[NUM_STRENGTHS];
  float maxima[NUM_STRENGTHS];
  // long fft_len = client.ap_shmem->ap_gdata.state.fft_len;
  int fft_len2 = 1024;
  // int fft_len2 = 32768;
  long datasize = 8 * 1024 * 1024;
  FILE* outfile = fopen("temp.txt", "w");
  int num_ffts = 0;

  for(int i = 0; i < NUM_STRENGTHS; i++) {
    totals[i] = 0;
    maxima[i] = 0;
  }

  // Note rawdata is in BYTES, each byte is 4 samples, fft_len is in samples, so # bytes = fft_len2 / 4
  // data_chunk_now is also in bytes, but it INCREMENTS by fft_len2 /2 / 4 to produce overlap.
  // That is, it only increments by HALF of a FFT length each iteration.
  // To find # of samples from data_chunk_now, MULTIPLY BY 4!
  for(long data_chunk_now = 0; data_chunk_now < datasize - fft_len2/2/4; data_chunk_now += fft_len2/2/4) {
    num_ffts++;
    if (splitter_bits_to_float(&client.rawdata[data_chunk_now],
                             (float *)data, fft_len2)) {
      exit(SBTF_ERR);
    }
    fftwf_execute(fp2); // Note that this fft knows about fft_len2, and computes in chunks of fft_len2
			   // The maximum possible result is 1,024, average should be such that power = 1024 
    for(int i = 0; i < NUM_STRENGTHS; i++) { 
      strengths[i] = 0;
      long amt_to_subtract = (i==0)?0:pow_int(2,i-1);
      for(long j = 0; j < pow_int(2,i); j++) {
        float real = data[(j - amt_to_subtract + fft_len2) % fft_len2][0];
        float imag = data[(j - amt_to_subtract + fft_len2) % fft_len2][1];
        strengths[i] += real * real + imag * imag;
      }
      totals[i] += strengths[i];
      maxima[i] = (strengths[i] > maxima[i])?strengths[i]:maxima[i];
      // fprintf(outfile, "strengths[%d] = %f\n", i, strengths[i]);
      // fprintf(outfile, "totals[%d] = %f\n", i, totals[i]);
      // fprintf(outfile, "maxima[%d] = %f\n", i, maxima[i]);
      // fprintf(outfile, "strengths[%d] / %d = %f\n", i, pow_int(2, i), strengths[i] / pow(2, i));
      fflush(outfile);
    }
  }

  fprintf(outfile, "The averages below represent power per DC bin.\n");
  for(int i = 0; i < NUM_STRENGTHS; i++) {
    averages[i] = totals[i] / (float) num_ffts / (float) pow_int(2,i) / (float) fft_len2 / 2.0;
    fprintf(outfile, "averages[%d] = %f\n", i, averages[i]);
  } 
  for(int i = 0; i < NUM_STRENGTHS; i++) {
    fprintf(outfile, "normalized maxima[%d] = %f\n", i, maxima[i] / (float) pow_int(2,i) / (float) fft_len2 / 2.0);
  } 
  fflush(outfile);
}

int Science::get_indices_to_randomize(std::vector<long>& indices, int scale) {

  FILE* outfile;
  int init_indices = 100;
  int num_indices = 0;
  int size = init_indices;
  indices.resize(init_indices);
  if(scale > 7) {
      fprintf(stderr, "In Science::get_indices_to_randomize: error: scale > 7\n");
      printf("In Science::get_indices_to_randomize: error: scale > 7\n"); 
      exit(-1);
  }
  // power_threshold[scale]
  float power_threshold[] = { 12.1, 6.9, 5.3, 3.0, 2.46, 1.96, 1.62, 1.37};
  
  int fft_len2 = 1024;
  long datasize = 8 * 1024 * 1024;
  float strength;

  outfile = fopen("indices.txt", "w");

  for(long data_chunk_now = 0; data_chunk_now < datasize - fft_len2/2/4; data_chunk_now += fft_len2/2/4) {
    if (splitter_bits_to_float(&client.rawdata[data_chunk_now],
                             (float *)data, fft_len2)) {
      exit(SBTF_ERR);
    }
    fftwf_execute(fp2); // Note that this fft knows about fft_len2 (if you define it in astropulse.h), and computes in chunks of fft_len2
			   // The maximum possible result is 1,024, average should be such that power = 1024 
    
    /* For debugging: check power in DC bin for this data chunk
    if(data_chunk_now == 1359872) {
        printf("DC bin: %f %f, DC power: %f\n", data[0][0], data[0][1], (data[0][0] * data[0][0] + data[0][1] * data[0][1]) / fft_len2 / 2.0);
    }*/

    int expscale = pow_int(2, scale);
    strength = 0;
    for(int j = - expscale / 2; j < (expscale + 1) / 2; j++) {
      int j2 = (j + fft_len2) % fft_len2; 
      float real = data[j2][0];
      float imag = data[j2][1];
      strength += (real * real + imag * imag) / fft_len2 / 2.0 / expscale;
      // strength += (real * real + imag * imag) / fft_len2 / expscale;
    }

    /* For debugging: check that total power / 2.0 = fft_len2
    if(data_chunk_now < 10000) {
    float total = 0;
    for(int i = 0; i < fft_len2; i++) {
        total += (data[i][0] * data[i][0] + data[i][1] * data[i][1]) / fft_len2 / 2.0;    
    }
    printf("total = %f\n", total);
    }
    */

    if(strength > power_threshold[scale]) {
      fprintf(outfile, "At sample %d, data_chunk_now %d, strength %f > %f\n", data_chunk_now * 4, data_chunk_now, strength, power_threshold[scale]);
      indices[num_indices++] = data_chunk_now;
      if(num_indices >= size)
        size = num_indices + init_indices;
        indices.resize(size);
    }
  }

  fclose(outfile);
  return num_indices;
}

// Read in as many data chunks as needed from client.rawdata.
// Warning: this function messes with the Science data array.
void Science::generate_envelope(const std::vector<long>& indices, int num_indices) {

  // initialize the envelope array
  for(int i=0; i<fftlen_envelope; i++)
     envelope[i] = 0;

  // GOAL: find envelope (as it should be after one-bitting)
 
  // For counting to see if we've performed enough ffts
  // for a meaningful power envelope.
  int num_ffts_performed = 0;

  // Check whether envelope size is (wrongly) larger than fft len
  // This should never happen, but we check just in case.
  if(fftlen_envelope > client.ap_shmem->ap_gdata.state.fft_len) {
    printf("Error in ap_remove_radar.cpp : Science::generate_envelope : fftlen_envelope is > client fft_len and should not be.\n");
    fprintf(stderr, "Error in ap_remove_radar.cpp : Science::generate_envelope : fftlen_envelope is > client fft_len and should not be.\n");
  }

  // * cycle through the #s in the data, adding each of them to
  //   envelope_part
  //   env_index tells which envelope_part we are on
  bool good_data_chunk;
  int data_chunk_size = client.ap_shmem->ap_gdata.state.fft_len / 4;
  for(int data_chunk=0; 
          data_chunk <= client.rawdata.size() - data_chunk_size; 
          data_chunk+= data_chunk_size) {

    // Check whether this is a bad data chunk (within 100000 bytes of RFI.) 
    good_data_chunk = true;
    for(int n=0; n<num_indices; n++) {
      if( indices[n] < data_chunk + 100000 && indices[n] > data_chunk - 100000) {
        good_data_chunk = false;
	break;
      }
    }
    if(!good_data_chunk) continue;

    // Extract our data chunk into the data array.
    splitter_bits_to_float(&client.rawdata[data_chunk], 
		     	   (float *)data, 
			   client.ap_shmem->ap_gdata.state.fft_len);
    for(int env_index=0;
            env_index<client.ap_shmem->ap_gdata.state.fft_len / fftlen_envelope;
            env_index++) {
      for(int i = 0; i < fftlen_envelope; i++) {
        int d = env_index * fftlen_envelope + i;

        envelope_part[i][0] = data[d][0]; 
        envelope_part[i][1] = data[d][1];
      }
      fftwf_execute(fp_envelope);
      for(int i = 0; i < fftlen_envelope; i++) {
        envelope[i] += envelope_part[i][0] * envelope_part[i][0]
                     + envelope_part[i][1] * envelope_part[i][1];
      }
      num_ffts_performed++;
    }
  }

  // * if you don't get very many envelope_part's, print an error msg and exit.
  //   (this means there was too much RFI.)
  int min_ffts = 100;
  if(num_ffts_performed < min_ffts) {
    printf("Error in ap_remove_radar.cpp : generate_envelope: num_ffts_performed < %d.  Blanking too much RFI?\n", min_ffts);
    fprintf(stderr, "Error in ap_remove_radar.cpp: generate_envelope: num_ffts_performed < %d.  Blanking too much RFI?\n", min_ffts);
    // Quit early
    fflush(stdout);
    client.outfile.write_main(client.signal_vector);
    client.outfile.finish_results();
    client.finish();
    main_timer.stop();
    boinc_finish(0);
    exit(0);
  }

  // * when you're done, normalize the power array:
  //    * add up the powers, and divide by fftlen_envelope
  //      to get the average power
  //    * divide the whole array by the average power
  //    * array is avp(nu), the average power at frequency nu, normalized
  //      so that <avp(nu)>_nu = 1.
  double total_power = 0;
  for(int i = 0; i < fftlen_envelope; i++)
    total_power += envelope[i];
  double avg_power = total_power / fftlen_envelope;
  for(int i = 0; i < fftlen_envelope; i++)
    envelope[i] /= avg_power;
}

void Science::generate_pre_envelope() {

  double known_alpha = 0.547539;
  
  double total_power = 0;
  for(int i = 0; i < fftlen_envelope; i++) {
    pre_envelope[i] = (1 + known_alpha) * envelope[i] - known_alpha;
    if(pre_envelope[i] < 0) pre_envelope[i] = 0;
    total_power += pre_envelope[i];
  }
  if(total_power <= 0) {
     printf("Error in ap_remove_radar.cpp: generate_pre_envelope: total <= 0.\n");
     fprintf(stderr, "Error in ap_remove_radar.cpp: generate_pre_envelope: total <= 0.\n");
     exit(1);
  }
  for(int i = 0; i < fftlen_envelope; i++) // set average power per bin to 1.
    pre_envelope[i] *= (fftlen_envelope / total_power);
}

void Science::randomize_indices(const std::vector<long>& indices, int num_indices, int data_chunk_now, int fft_len, state_t& state) { 

// FILE* randsampfile = fopen("randsamp.txt", "a");
long prev_index = data_chunk_now - 1;
long prev_sample = (data_chunk_now - 1) * 4;
long this_index;
long this_sample;
   
    /* For debugging: print values of bits before they are randomized. 
    if(data_chunk_now == 1359872 - 4096) {
        printf("data_chunk_now: 1359872 - 4096\n");
        for(int k = 16384; k < 16384 + 1024; k++) {
            printf("%f %f | ", data[k][0], data[k][1]);
            if(k % 16 == 0) printf("\n");
        }
    }
    */
  
  // printf("data_chunk_now = %d\n", data_chunk_now); // for debugging
  for(int current = 0; current < num_indices; current++) {
    // If we're within 100000 bytes of a DC pulse, randomize this data chunk.
    if(indices[current] < data_chunk_now + 100000 && indices[current] > data_chunk_now - 100000) {
      // Note that we WILL randomize the same data twice, but that won't count twice toward numSamplesBlanked:
      // if we just blanked the last data chunk, then this data chunk only
      // adds fft_len / 2 samples, because of the overlap.
      if(state.lastDataChunkBlanked == data_chunk_now - fft_len / 4 / 2)
        state.numSamplesBlankedNew += fft_len / 2;
      else
        state.numSamplesBlankedNew += fft_len;
      double fractionBlanked = (double) state.numSamplesBlankedNew / NUM_SAMPLES;
      state.lastDataChunkBlankedNew = data_chunk_now;
      long lastSampleBlanked = state.lastDataChunkBlankedNew * 4L + fft_len - 1;
      int debug = 0;
      randomize_samples(fft_len, state.mtr_new, debug);
      // fprintf(randsampfile, "Num samples blanked: %ld, fraction blanked: %f, last sample blanked: %ld\n", state.numSamplesBlankedNew, 
      //       fractionBlanked, lastSampleBlanked);
      break;
    }
  }
  // fclose(randsampfile);
}


double gasdev(MtRand &mtr) {
  static int iset=0;
  static double gset;
  double fac, rsq, v1, v2;
  if(iset == 0) {
    do {
      v1=2.0*mtr.drand32()-1.0;
      v2=2.0*mtr.drand32()-1.0;
      rsq=v1*v1+v2*v2;
    } while(rsq >= 1.0 || rsq==0.0);
    fac=sqrt(-2.0*log(rsq)/rsq);
    gset=v1*fac;
    iset=1;
    return v2*fac;
  } else {
    iset=0;
    return gset;
  }
}

// One-bits the data (called "noise") which has length fftlen_noise
// That is, values greater than 0 are set to 1, and less than 0 are
// set to -1.
void one_bit(fftwf_complex * noise, int fftlen_noise) {
  int sum_0_0 = 0;
  int sum_0_1 = 0;
  for(int k=0; k<fftlen_noise; k++) {
    noise[k][0] = (noise[k][0] < 0) ? -1 : 1;
    noise[k][1] = (noise[k][1] < 0) ? -1 : 1;
    sum_0_0 += noise[k][0];
    sum_0_1 += noise[k][1];
  }
  // printf("noise[0][0] = %f, noise[0][1] = %f\n", noise[0][0], noise[0][1]);
  // printf("sum_0_0 = %d, sum_0_1 = %d\n", sum_0_0, sum_0_1);
}

// This generates unnormalized time domain noise with frequency spectrum given by "envelope".
// (Because it will be one-bitted immediately afterwards.)
void generate_time_noise_with_envelope(fftwf_complex * noise, double * envelope, int fftlen_noise, int fftlen_envelope, fftwf_plan fp_noise_inv, MtRand &mtr) {

  // double rnd1;
  // double rnd2;
  // double HALF_RAND_MAX = RAND_MAX/2.0;
  double multiplier;

  //  * generate fftlen_patch random cplx bit pairs (as doubles)
  /*
  for(int k=0; k<fftlen_noise; k++) {
    rnd1=static_cast<double>(rand());
    rnd1=rnd1 <= HALF_RAND_MAX ? -1 : 1;
    rnd2=static_cast<double>(rand());
    rnd2=rnd2 <= HALF_RAND_MAX ? -1 : 1;
    noise[k][0]=rnd1;
    noise[k][1]=rnd2;
  }*/
  //  * fourier transform them to get a Gaussian distribution
  //      average power per frequency bin is some constant P, and
  //      expected power in each bin is <P>.
  // fftwf_execute(fp_noise);
  // avg power per bin = 2 * fftlen_patch.
  //  * TODO: divide by sqrt(P), so avg power in each bin is 1.
  /* double sqrt_avg_power = sqrt((double) 2 * fftlen_noise);
  for(int k = 0; k<fftlen_noise; k++) {
    noise[k][0] /= sqrt_avg_power;
    noise[k][1] /= sqrt_avg_power;
  } */
  for(int k=0; k<fftlen_noise; k++) {
    noise[k][0] = gasdev(mtr);
    noise[k][1] = gasdev(mtr);
  }
  //  * multiply patch by sqrt(avp(nu)).
  //      expected power in bin nu is now avp(nu)
  for(int k = 0; k<fftlen_noise; k++) {
    // 0.5 is added so we get 0.5, ... fftlen_noise - 0.5, which is then rounded down.
    int index = (int) ( ((float)k + 0.5) / (float)fftlen_noise * (float)fftlen_envelope);
    multiplier = sqrt(envelope[ index ]);

    noise[k][0] *= multiplier;
    noise[k][1] *= multiplier;
    // if(k == 0)
    //  printf("noise[0][0] = %f, noise[0][1] = %f, noise power = %f, multiplier = %f, index = %d\n", noise[0][0], noise[0][1], noise[0][1] * noise[0][1] + noise[0][0] * noise[0][0], multiplier, index);
  }
  //  * inverse FFT
  fftwf_execute(fp_noise_inv);
}



// Use debug flag if you want randomize_samples to print something out
// on specific samples, for instance.
void Science::randomize_samples(int fft_len, MtRand& mtr, int debug) {
  // FILE *randsampfile;
  // randsampfile = fopen("randsamp.txt", "a");
  // fprintf(randsampfile, "in randomize_samples: in: %s\n", mtr.toString());

  int patch_index = 0;

  for(int index = 0; index < fft_len; index++) {
    if(patch_index == 0) {
      generate_time_noise_with_envelope(patch, pre_envelope, fftlen_patch, fftlen_envelope, fp_patch_inv, mtr);
      one_bit(patch, fftlen_patch);
    }
// * use bits from this distribution to fill the holes
    while(index < client.ap_shmem->ap_gdata.state.fft_len) {
      data[index][0] = patch[patch_index][0];
      data[index][1] = patch[patch_index][1];
    // data[index][0] = mtr.drand32() > 0.5 ? 1.0 : - 1.0; 
    // data[index][1] = mtr.drand32() > 0.5 ? 1.0 : - 1.0; 
      patch_index++;
      if(patch_index >= fftlen_patch) {
// * when you don't have any patch bits left,
// start a new patch, going back through the for loop
        patch_index = 0;
        break;
      }
      index++;
    }
    patch_index = 0;
  } // end "for(int index = 0; index < fft_len; index++)"

  // fprintf(randsampfile, "in randomize_samples: out: %s\n", mtr.toString());
  // fflush(randsampfile);
  // fclose(randsampfile);
}

} // end namespace Astropulse
