<commit_before>extern crate dft;

#[derive(Clone, Debug)]
pub struct FrequencyBucket {
    pub min_freq: f64,
    pub max_freq: f64,
    pub intensity: f64
}

impl FrequencyBucket {
    pub fn ave_freq(&self) -> f64 {
        (self.min_freq + self.max_freq) / 2.0
    }
}

pub fn fft(input: &Vec<f64>, sample_rate: f64) -> Vec<FrequencyBucket> {
    let frames = input.len();
    let mean_input = input.iter().sum::<f64>()/input.len() as f64;
    
    let mut intensities = input.iter().map(|x|x-mean_input).collect::<Vec<_>>();
    let plan = dft::Plan::new(dft::Operation::Forward, frames);
    dft::transform(&mut intensities, &plan);

    let frequency_resolution = sample_rate / 2.0 / frames as f64;
    
    intensities.iter().enumerate().map(|(index, &value)| {
        let index = index as f64;
        FrequencyBucket {
            min_freq: index * frequency_resolution,
            max_freq: (index+1.0) * frequency_resolution,
            intensity: value
        }
    }).collect()
}

pub fn find_fundamental_frequency(frequency_domain: &Vec<FrequencyBucket>) -> Option<f64> {
    //TODO look at all significant frequencies, find fundamental
    //TODO return None is none of them are significant
    let positive_buckets = frequency_domain.iter().filter(|x| x.intensity > 0.0).cloned().collect::<Vec<_>>();
    let average_intensity = positive_buckets.iter().map(|x| x.intensity).sum::<f64>() / frequency_domain.len() as f64;
    let significant_buckets = positive_buckets.iter().filter(|x| x.intensity > average_intensity).cloned().collect::<Vec<_>>();
    
    let max_bucket = significant_buckets.iter()
        .fold(None as Option<::transforms::FrequencyBucket>, |max, next|
              if max.is_none() || max.clone().unwrap().intensity < next.intensity { Some(next.clone()) } else { max }
        ).unwrap();

    Some(max_bucket.ave_freq())
}

pub fn correlation(input: &Vec<f64>) -> Vec<f64> {
    let mut correlation = Vec::with_capacity(input.len());
    for offset in 0..input.len() {
        let mut c = 0.0;
        for i in 0..input.len()-offset {
            let j = i+offset;
            c += input[i] * input[j];
        }
        correlation.push(c);
    }
    correlation
}

pub fn find_fundamental_frequency_correlation(input: &Vec<f64>, sample_rate: f64) -> Option<f64> {
    let correlation = correlation(&input);

    let mut first_peak_width = 0;
    for offset in 0..correlation.len() {
        if correlation[offset] < 0.0 {
            first_peak_width = offset;
            break;
        }
    }
    if first_peak_width == 0 {
        return None;
    }

    let peak = correlation.iter()
        .enumerate()
        .skip(first_peak_width)
        .fold((first_peak_width, 0.0 as f64), |(xi, xmag), (yi, &ymag)| if ymag > xmag { (yi, ymag) } else { (xi, xmag) });

    let (peak_index, _) = peak;
    Some(sample_rate / peak_index as f64)
}

fn interpolate(correlation: &Vec<f64>, x: f64) -> f64 {
    if x < 0.0 {
        println!("<0");
        correlation[0]
    }
    else if x >= correlation.len() as f64 {
        println!(">len");
        correlation[correlation.len()-1]
    }
    else {
        let x0 = x.floor();
        let y0 = correlation[x0 as usize];
        let x1 = x.ceil();
        let y1 = correlation[x1 as usize];

        (y0*(x1-x) + y1*(x-x0)) / (x1-x0)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64::consts::PI;
    
    const SAMPLE_RATE: f64 = 44100.0;
    const FRAMES: usize = 512;

    fn frequency_resolution() -> f64 {
        SAMPLE_RATE / 2.0 / FRAMES as f64
    }

    fn sin_arg(f: f64, t: f64, phase: f64) -> f64 {
        2.0 as f64 * PI * f * t + phase
    }
<commit_msg>Added binary search to refine fundamental frequency search on the correlation<commit_after>extern crate dft;

#[derive(Clone, Debug)]
pub struct FrequencyBucket {
    pub min_freq: f64,
    pub max_freq: f64,
    pub intensity: f64
}

impl FrequencyBucket {
    pub fn ave_freq(&self) -> f64 {
        (self.min_freq + self.max_freq) / 2.0
    }
}

pub fn fft(input: &Vec<f64>, sample_rate: f64) -> Vec<FrequencyBucket> {
    let frames = input.len();
    let mean_input = input.iter().sum::<f64>()/input.len() as f64;
    
    let mut intensities = input.iter().map(|x|x-mean_input).collect::<Vec<_>>();
    let plan = dft::Plan::new(dft::Operation::Forward, frames);
    dft::transform(&mut intensities, &plan);

    let frequency_resolution = sample_rate / 2.0 / frames as f64;
    
    intensities.iter().enumerate().map(|(index, &value)| {
        let index = index as f64;
        FrequencyBucket {
            min_freq: index * frequency_resolution,
            max_freq: (index+1.0) * frequency_resolution,
            intensity: value
        }
    }).collect()
}

pub fn find_fundamental_frequency(frequency_domain: &Vec<FrequencyBucket>) -> Option<f64> {
    //TODO look at all significant frequencies, find fundamental
    //TODO return None is none of them are significant
    let positive_buckets = frequency_domain.iter().filter(|x| x.intensity > 0.0).cloned().collect::<Vec<_>>();
    let average_intensity = positive_buckets.iter().map(|x| x.intensity).sum::<f64>() / frequency_domain.len() as f64;
    let significant_buckets = positive_buckets.iter().filter(|x| x.intensity > average_intensity).cloned().collect::<Vec<_>>();
    
    let max_bucket = significant_buckets.iter()
        .fold(None as Option<::transforms::FrequencyBucket>, |max, next|
              if max.is_none() || max.clone().unwrap().intensity < next.intensity { Some(next.clone()) } else { max }
        ).unwrap();

    Some(max_bucket.ave_freq())
}

pub fn correlation(input: &Vec<f64>) -> Vec<f64> {
    let mut correlation = Vec::with_capacity(input.len());
    for offset in 0..input.len() {
        let mut c = 0.0;
        for i in 0..input.len()-offset {
            let j = i+offset;
            c += input[i] * input[j];
        }
        correlation.push(c);
    }
    correlation
}

pub fn find_fundamental_frequency_correlation(input: &Vec<f64>, sample_rate: f64) -> Option<f64> {
    let correlation = correlation(&input);

    let mut first_peak_width = 0;
    for offset in 0..correlation.len() {
        if correlation[offset] < 0.0 {
            first_peak_width = offset;
            break;
        }
    }
    if first_peak_width == 0 {
        return None;
    }

    let peak = correlation.iter()
        .enumerate()
        .skip(first_peak_width)
        .fold((first_peak_width, 0.0 as f64), |(xi, xmag), (yi, &ymag)| if ymag > xmag { (yi, ymag) } else { (xi, xmag) });

    let (peak_index, _) = peak;

    let refined_peak_index = refine_fundamentals(&correlation, peak_index as f64);
    
    Some(sample_rate / refined_peak_index)
}

fn refine_fundamentals(correlation: &Vec<f64>, initial_guess: f64) -> f64 {
    let mut low_bound = initial_guess - 0.5;
    let mut high_bound = initial_guess + 0.5;

    for _ in 0..5 {
        let data_points = 2 * correlation.len() / high_bound.ceil() as usize;
        let low_guess = score_guess(&correlation, low_bound, data_points);
        let high_guess = score_guess(&correlation, high_bound, data_points);

        let midpoint = (low_bound + high_bound) / 2.0;
        if high_guess > low_guess {
            low_bound = midpoint;
        }
        else {
            high_bound = midpoint;
        }
    }
    (low_bound + high_bound) / 2.0
}

fn score_guess(correlation: &Vec<f64>, period: f64, data_points: usize) -> f64 {
    let mut score = 0.0;
    for i in 0..data_points {
        if i % 2 == 0 {
            score += i as f64 * interpolate(&correlation, i as f64 * period / 2.0);
        }
        else {
            score -= i as f64 * interpolate(&correlation, i as f64 * period / 2.0);
        }
    }
    score
}

fn interpolate(correlation: &Vec<f64>, x: f64) -> f64 {
    if x < 0.0 {
        println!("<0");
        correlation[0]
    }
    else if x >= correlation.len() as f64 {
        println!(">len");
        correlation[correlation.len()-1]
    }
    else {
        let x0 = x.floor();
        let y0 = correlation[x0 as usize];
        let x1 = x.ceil();
        let y1 = correlation[x1 as usize];

        if (x0 as usize == x1 as usize) {
            y0
        }
        else {
            (y0*(x1-x) + y1*(x-x0)) / (x1-x0)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64::consts::PI;
    
    const SAMPLE_RATE: f64 = 44100.0;
    const FRAMES: usize = 512;

    fn frequency_resolution() -> f64 {
        SAMPLE_RATE / 2.0 / FRAMES as f64
    }

    fn sin_arg(f: f64, t: f64, phase: f64) -> f64 {
        2.0 as f64 * PI * f * t + phase
    }

{'fraction_non_alphanumeric': 0.14452402957486138, 'fraction_numerical': 0.03153881700554528, 'mean_word_length': 2.4531312325488632, 'pattern_counts': {'":': 0, '<': 49, '<?xml version=': 0, '>': 70, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1602234', 'n_tokens_mistral': 3123, 'n_tokens_neox': 2895, 'n_words': 707}