use num::complex::Complex;
use std::f32;
use std::fs::File;
use std::io::BufWriter;
use std::path::Path;
use hound;
use png::HasParameters;
use crate::colour_gradient::{ColourGradient, RGBAColour};
use crate::errors::SonogramError;
use crate::utility;
type Spectrogram = Vec<Vec<Complex<f32>>>;
type WindowFn = fn(u32, u32) -> f32;
pub struct SpecOptionsBuilder {
width: u32,
height: u32,
sample_rate: u32,
data: Vec<i16>,
window: WindowFn,
greyscale: bool,
verbose: bool,
}
impl SpecOptionsBuilder {
pub fn new(width: u32, height: u32) -> Self {
SpecOptionsBuilder {
width,
height,
sample_rate: 8000,
data: vec![],
window: utility::rectangular,
greyscale: false,
verbose: false,
}
}
pub fn set_window_fn(&mut self, window: fn(u32, u32) -> f32) -> &mut Self {
self.window = window;
self
}
pub fn load_data_from_file(&mut self, fname: &Path) -> Result<&mut Self, SonogramError> {
let mut reader = hound::WavReader::open(fname)?;
assert_eq!(reader.spec().bits_per_sample, 16);
assert_eq!(reader.spec().channels, 1);
let data = reader.samples().map(|x| x.unwrap()).collect();
let sample_rate = reader.spec().sample_rate;
Ok(self.load_data_from_memory(data, sample_rate))
}
pub fn load_data_from_memory(&mut self, data: Vec<i16>, sample_rate: u32) -> &mut Self {
self.data = data;
self.sample_rate = sample_rate;
self
}
pub fn downsample(&mut self, divisor: usize) -> &mut Self {
if divisor == 0 {
panic!("The divisor is too small");
}
if divisor == 1 {
return self;
}
if self.data.is_empty() {
panic!("Need to load the data before calling this function");
}
let mut j = 0;
for i in (0..self.data.len() - divisor).step_by(divisor) {
let sum: i32 = self.data[i..i + divisor]
.iter()
.fold(0i32, |mut sum, &val| {
sum += i32::from(val);
sum
});
let avg = sum / divisor as i32;
self.data[j] = avg as i16;
j += 1;
}
self.data.resize(self.data.len() / divisor, 0);
self.sample_rate /= divisor as u32;
self
}
pub fn set_verbose(&mut self) -> &mut Self {
self.verbose = true;
self
}
pub fn set_greyscale(&mut self) -> &mut Self {
self.greyscale = true;
self
}
pub fn build(&self) -> Spectrograph {
if self.data.is_empty() {
panic!("SpecOptionsBuilder requires data to be loaded")
}
let audio_length_sec = self.data.len() as u32 / self.sample_rate;
if self.verbose {
println!("Length (s): {}", audio_length_sec);
}
let mut gradient = ColourGradient::new();
if self.greyscale {
gradient.add_colour(RGBAColour::new(0, 0, 0, 0));
gradient.add_colour(RGBAColour::new(255, 255, 255, 0));
} else {
gradient.add_colour(RGBAColour::new(0, 0, 0, 0));
gradient.add_colour(RGBAColour::new(55, 0, 110, 0));
gradient.add_colour(RGBAColour::new(0, 0, 180, 0));
gradient.add_colour(RGBAColour::new(0, 255, 255, 0));
gradient.add_colour(RGBAColour::new(0, 255, 0, 0));
gradient.add_colour(RGBAColour::new(255, 255, 0, 0));
gradient.add_colour(RGBAColour::new(230, 160, 0, 0));
gradient.add_colour(RGBAColour::new(255, 0, 0, 0));
}
Spectrograph {
width: self.width,
height: self.height,
data: self.data.clone(),
window: self.window,
spectrogram: vec![vec![]],
gradient,
verbose: self.verbose,
}
}
}
pub struct Spectrograph {
width: u32,
height: u32,
data: Vec<i16>,
window: WindowFn,
spectrogram: Spectrogram,
gradient: ColourGradient,
verbose: bool,
}
impl Spectrograph {
pub fn omega(&self, p: f32, q: f32) -> Complex<f32> {
let trig_arg = 2.0 * f32::consts::PI * q / p;
Complex::new(f32::cos(trig_arg), f32::sin(trig_arg))
}
pub fn compute(&mut self, chunk_len: usize, overlap: f32) {
assert!(0.0 <= overlap && overlap < 1.0);
let step = (chunk_len as f32 * (1.0 - overlap)) as usize;
if self.verbose {
println!("Computing spectrogram...");
println!(
"Chunk: {} Overlap: {}",
chunk_len,
overlap * chunk_len as f32
);
println!("Step len: {}", step);
println!("Data len: {}", self.data.len());
}
let mut new_len = 0;
while new_len + chunk_len < self.data.len() {
new_len += step;
}
if new_len != self.data.len() {
new_len += chunk_len;
let padding = &mut vec![0; new_len - self.data.len()];
self.data.append(padding);
}
self.chunkify(chunk_len, step);
}
pub fn save_as_png(&mut self, fname: &Path, log_mode: bool) -> Result<(), std::io::Error> {
let data_len = self.spectrogram[0].len();
let multiplier = 0.5;
let img_len_used = data_len as f32 * multiplier;
let log_coef = 1.0 / (self.height as f32 + 1.0).log(f32::consts::E) * img_len_used;
let file = File::create(fname)?;
let w = &mut BufWriter::new(file);
let mut encoder = png::Encoder::new(w, self.width, self.height);
encoder.set(png::ColorType::RGBA).set(png::BitDepth::Eight);
let mut writer = encoder.write_header()?;
let mut img: Vec<u8> = vec![];
for y in (0..self.height).rev() {
for x in 0..self.width {
let freq = if log_mode {
img_len_used - (log_coef * (self.height as f32 + 1.0 - y as f32).log(f32::consts::E))
} else {
let ratio = y as f32 / self.height as f32;
ratio * img_len_used
};
let colour = self.get_colour(self.spectrogram[x as usize][freq as usize], 15.0);
img.extend(colour.to_vec());
}
}
writer.write_image_data(&img)?;
Ok(())
}
fn get_colour(&mut self, c: Complex<f32>, threshold: f32) -> RGBAColour {
let value = 0.5 * (c.norm_sqr() + 1.0).log10();
self.gradient.set_max(threshold);
self.gradient.get_colour(value).clone()
}
fn get_number_of_chunks(&mut self, chunk_len: usize, step: usize) -> usize {
let mut i = 0;
let mut chunks = 0;
while i + chunk_len <= self.data.len() {
i += step;
chunks += 1;
}
if i == self.data.len() {
chunks -= 1;
}
chunks
}
fn chunkify(&mut self, chunk_len: usize, step: usize) {
self.spectrogram.clear();
self.spectrogram.reserve(self.width as usize);
let num_chunks = self.get_number_of_chunks(chunk_len, step) as f32;
if self.verbose {
println!("Number of Chunks: {}", num_chunks);
}
let chunk_width_ratio = num_chunks / self.width as f32;
let mut j = 0.0;
while j < num_chunks {
let p = j as usize * step;
let mut signal: Vec<Complex<f32>> = self.data[p..]
.iter()
.take(chunk_len)
.map(|d| Complex::new(f32::from(*d), 0.0))
.collect();
self.transform(&mut signal);
self.spectrogram.push(signal);
j += chunk_width_ratio;
}
}
fn transform(&mut self, signal: &mut Vec<Complex<f32>>) {
let min_len = signal.len();
let power = pad_to_power2(signal, min_len);
if power == 0 {
return;
}
let mut transformed = vec![Complex::new(0f32, 0f32); signal.len()];
for i in 0..signal.len() {
transformed[utility::reverse_bits(i, power)] =
signal[i] * (self.window)(i as u32, signal.len() as u32);
}
let mut n = 2;
while n <= transformed.len() {
let mut i = 0;
while i <= transformed.len() - n {
for m in i..(i + n / 2) {
let term1 = transformed[m];
let term2 = self.omega(n as f32, -(m as f32)) * transformed[m + n / 2];
transformed[m] = term1 + term2;
transformed[m + n / 2] = term1 - term2;
}
i += n;
}
n *= 2;
}
signal.clear();
signal.extend(transformed.into_iter());
}
}
fn pad_to_power2(signal: &mut Vec<Complex<f32>>, min_len: usize) -> usize {
let mut power = 1;
let mut new_len = 2;
while new_len < min_len {
new_len *= 2;
power += 1;
}
pad(signal, new_len);
let padding = &mut vec![Complex::new(0.0, 0.0); new_len - signal.len()];
signal.append(padding);
power
}
fn pad(signal: &mut Vec<Complex<f32>>, new_len: usize) {
if new_len > signal.len() {
signal.resize_with(new_len, || Complex::new(0.0, 0.0));
}
}