extern crate goertzel_filter;
extern crate num;
extern crate itertools_num;

use num::Complex;
use itertools_num::linspace;
use std::ops::Sub;
use std::f64::consts::{PI};

static TAU: f64 = 2.0 * PI;

/// Tests are derived from this paper excerpt:
/// http://dsp.stackexchange.com/a/635/5614

#[test]
fn dft_linearity() {
  // This test asserts that the DFT of a summed pair of signals at a given
  // frequency is the same as the sum of the DFTs of those signals taken
  // separately.

  // 440Hz as wave frequency (middle A)
  let freq: f64 = 440.0;
  let interval_length = 1.0;
  // Time vector sampled at 880 times/s (~Nyquist), over 1s
  let delta: f64 = &interval_length / &freq / 2.0;
  let time_1s = linspace(0.0, interval_length, (&freq / 2.0) as usize)
    .map(|sample| { sample * delta });

  
  let (sine_440,
       sine_100): (Vec<f64>, Vec<f64>) =
    time_1s.map(|time_sample| {
      let sample_440 = (freq * &time_sample).sin();
      let sample_100 = (100.0 * &time_sample).sin();
      
      (sample_440, sample_100)
    }).unzip();

  let summed_signal: Vec<f64> = sine_440
    .iter()
    .zip(sine_100.iter())
    .map(|(sample_440, sample_100)| {
      sample_440 + sample_100
    }).collect();
    
  let dft_440 = goertzel_filter::dft(&sine_440, 440.0);  
  let dft_100 = goertzel_filter::dft(&sine_100, 440.0);
  
  let summed_signal_dft = goertzel_filter::dft(&summed_signal, 440.0);
  
  let summed_dft = dft_440 + dft_100;

  let complex_diff = summed_dft - summed_signal_dft;
  
  assert!(complex_diff.re.abs() < 0.00001);
  assert!(complex_diff.im.abs() < 0.00001);
}

#[test]
fn dft_power_linearity() {
  // This test asserts that the sum of the signal powers of two different
  // signals at a given frequency is the same as the power of the sum of the
  // signals.

  // 440Hz as wave frequency (middle A)
  let freq: f64 = 440.0;
  let interval_length = 1.0;
  // Time vector sampled at 880 times/s (~Nyquist), over 1s
  let delta: f64 = &interval_length / &freq / 2.0;
  let time_1s = linspace(0.0, interval_length, (&freq / 2.0) as usize)
    .map(|sample| { sample * delta });

  
  let (sine_440,
       sine_100): (Vec<f64>, Vec<f64>) =
    time_1s.map(|time_sample| {
      let sample_440 = (freq * &time_sample).sin();
      let sample_100 = (100.0 * &time_sample).sin();
      
      (sample_440, sample_100)
    }).unzip();

  let summed_signal: Vec<f64> = sine_440
    .iter()
    .zip(sine_100.iter())
    .map(|(sample_440, sample_100)| {
      sample_440 + sample_100
    }).collect();
    
  let power_440 = goertzel_filter::dft_power(&sine_440, 440.0);  
  let power_100 = goertzel_filter::dft_power(&sine_100, 440.0);
  
  let summed_signal_power = goertzel_filter::dft_power(&summed_signal, 440.0);
  
  let summed_power = (power_440.sqrt() + power_100.sqrt()).powi(2);

  assert!((summed_power - summed_signal_power).abs() < 0.00001);
}

#[test]
fn dft_const() {
  // This test asserts that the DFT of an impulse at any frequency is constant.
  
  // Maximum frequency (so sample at double this)
  let max_freq: f64 = 500.0;
  let interval_length = 0.5;
  
  // Time vector sampled at 1000 times/s (~Nyquist), over 500ms
  let delta: f64 = &interval_length / &max_freq / 2.0;
  let time_500ms = linspace(0.0, interval_length, (&max_freq / 2.0) as usize)
    .map(|sample| { sample * delta });

  // Unit impulse at time zero
  let mut impulse_500ms: Vec<f64> = time_500ms.collect();
  impulse_500ms[0] = 1.0;

  let dft_bins: Vec<f64> = vec![1.0, 10.0, 100.0, 250.0, 300.0, 500.0];
  
  let impulse_dfts: Vec<Complex<f64>> = dft_bins.iter().map(|&dft_freq| {
    goertzel_filter::dft(&impulse_500ms, dft_freq)
  }).collect();

  for dft_pair in impulse_dfts.chunks(2) {
    let dft_diff = dft_pair[0].norm() - dft_pair[1].norm();
    assert!(dft_diff.abs() < 0.001);
  }
}


#[test]
fn dft_const_power() {
  // This test asserts that the DFT power of an impulse at any frequency is
  // constant.
  
  // Maximum frequency (so sample at double this)
  let max_freq: f64 = 500.0;
  let interval_length = 0.5;
  
  // Time vector sampled at 1000 times/s (~Nyquist), over 500ms
  let delta: f64 = &interval_length / &max_freq / 2.0;
  let time_500ms = linspace(0.0, interval_length, (&max_freq / 2.0) as usize)
    .map(|sample| { sample * delta });

  // Unit impulse at time zero
  let mut impulse_500ms: Vec<f64> = time_500ms.collect();
  impulse_500ms[0] = 1.0;

  let dft_bins: Vec<f64> = vec![1.0, 10.0, 100.0, 250.0, 300.0, 500.0];
  
  let impulse_dfts: Vec<f64> = dft_bins.iter().map(|&dft_freq| {
    goertzel_filter::dft_power(&impulse_500ms, dft_freq)
  }).collect();

  for pow_pair in impulse_dfts.chunks(2) {
    assert!((pow_pair[0] - pow_pair[1]).abs() < 0.001);
  }
}

#[test]
fn dft_const_phase_difference() {
  // This test asserts that the phase of the DFT of a bunch of time-offset
  // impulses is linear (in other words, that their first difference is roughly
  // constant) modulo tau.
  
  // Maximum frequency (so sample at double this)
  let max_freq: f64 = 500.0;
  let interval_length = 0.5;
  
  // Time vector sampled at 1000 times/s (~Nyquist), over 500ms
  let delta: f64 = &interval_length / &max_freq / 2.0;
  let time_500ms: Vec<f64> = linspace(0.0,
                                      interval_length,
                                      (&max_freq / 2.0) as usize)
    .map(|sample| { sample * delta })
    .collect();

  // Number of offsets to use
  let offset_count = 15;
  // How far apart the offset is in sample if they were equally spaced
  let offset_length = (time_500ms.len() / &offset_count) as usize;

  // Create a bunch of unit impulses offset by that many samples
  let offset_impulses: Vec<Vec<f64>> = (0..offset_count)
    .map(|offset_ind| {
      let mut impulse: Vec<f64> = time_500ms.clone();

      impulse[offset_ind as usize * offset_length] = 1.0;
      
      impulse
    })
    .collect();

  let dft_freq = 500.0;

  let impulse_dfts: Vec<Complex<f64>> = offset_impulses.iter().map(|impulse| {
    goertzel_filter::dft(&impulse, dft_freq)
  }).collect();

  let dft_phases: Vec<f64> = impulse_dfts
    .iter()
    .map(|&dft| { ((dft.arg() % TAU) + TAU) % TAU }) // modulo workaround
    .collect();

  let dft_phase_diffs: Vec<f64> = first_difference(&dft_phases)
    .iter().map(|&diff| { (diff + TAU ) % TAU }).collect();

  for phase_pair in dft_phase_diffs.chunks(2) {
    assert!((phase_pair[0] - phase_pair[1]).abs() < 0.001);
  }
}

// with great help from http://codereview.stackexchange.com/a/144037/105623
fn first_difference<T>(input: &[T]) -> Vec<T>
  where for <'a> &'a T: Sub<Output = T> {

  let minuend_iter = input.iter();
  let subtrahend_iter = input.iter().skip(1);
  
  minuend_iter.zip(subtrahend_iter).map(|(minuend, subtrahend)| {
    minuend - subtrahend
  }).collect()
}