// Generate Blue Noise Mask tables.

extern crate rayon;

use std::cmp::Ordering;
use std::env;
use std::fs::File;
use std::io::Write;
use std::ops::{Index, IndexMut};
use std::path::Path;

use rayon::prelude::*;

const WINDOW_RADIUS: isize = 63;
const FILTER_WIDTH: f32 = 1.2;
const FILTER_PASSES: usize = 1;

// These are specified in powers of two (2^N) for fast wrapping
// in the generated Rust code.
const NUM_MASKS_POW: usize = 7;  // 128
const MASK_SIZE_POW: usize = 7;  // 128

const MASK_SIZE: usize = 1 << MASK_SIZE_POW;
const MASK_SIZE_BITMASK: usize = (1 << MASK_SIZE_POW) - 1;

fn main() {
    let out_dir = env::var("OUT_DIR").unwrap();
    let dest_path = Path::new(&out_dir).join("blue_noise_masks.rs");
    let mut f = File::create(&dest_path).unwrap();

    // Generate masks
    let masks = (0..(1 << NUM_MASKS_POW))
        .collect::<Vec<_>>()
        .par_iter()
        .map(|i| blue_noise_mask(*i))
        .collect::<Vec<_>>();

    // Write the beginning bits of the file
    f.write_all(format!(r#"
// This file is automatically generated.

// Compute points of the Halton sequence with with Faure-permutations for different bases.

pub const NUM_MASKS: u32 = {};
pub const MASK_SIZE: u32 = {};
pub const NUM_MASKS_WRAP_BITMASK: u32 = {};
const MASK_SIZE_WRAP_BITMASK: u32 = {};
const MASK_POINTS: u32 = {};

#[inline]
pub fn get_point(x: u32, y: u32) -> &'static [f32] {{
    let i = get_index_to_point(x, y);

    &MASKS[i..(i + NUM_MASKS as usize)]
}}

#[inline]
pub fn get_index_to_point(x: u32, y: u32) -> usize {{
    let x = x & MASK_SIZE_WRAP_BITMASK;
    let y = y & MASK_SIZE_WRAP_BITMASK;

    ((y * MASK_SIZE * NUM_MASKS) + (x * NUM_MASKS)) as usize
}}
"#,
    1 << NUM_MASKS_POW,
    MASK_SIZE,
    (1 << NUM_MASKS_POW) - 1,
    MASK_SIZE_BITMASK,
    MASK_SIZE * MASK_SIZE,
    ).as_bytes()).unwrap();

    // Write the mask data
    f.write_all(format!(r#"
pub static MASKS: [f32; {}] = [
    "#, MASK_SIZE * MASK_SIZE * (1 << NUM_MASKS_POW)).as_bytes()).unwrap();

    for y in 0..MASK_SIZE {
        for x in 0..MASK_SIZE {
            for mask in masks.iter() {
                let v = mask[(x as isize, y as isize)];
                f.write_all(format!(r#"    {:.8},
"#, v).as_bytes()).unwrap();
            }
        }
    }

    f.write_all(format!(r#"
];
"#).as_bytes()).unwrap();
}



/// Creates a blue noise mask
fn blue_noise_mask(seed: u32) -> Mask {
    // Generate white noise mask
    let mut mask = Mask::new();
    for (i, v) in mask.data.iter_mut().enumerate() {
        *v = hash_u32_to_f32(i as u32, seed);
    }

    // High pass and remap
    for _ in 0..FILTER_PASSES {
        high_pass_filter(&mut mask, WINDOW_RADIUS, FILTER_WIDTH);
        remap_values(&mut mask);
    }

    mask
}

/// Performs a high pass filter on a Mask
fn high_pass_filter(mask: &mut Mask, window_radius: isize, filter_width: f32) {
    // Precompute filter convolution matrix
    let conv = {
        let mut conv = Mask::new();
        for j in (-window_radius)..(window_radius + 1) {
            for i in (-window_radius)..(window_radius + 1) {
                let n = (((j * j) + (i * i)) as f32).sqrt();
                //let n = (j.abs() + i.abs()) as f32;
                conv[(i + window_radius, j + window_radius)] = sinc(n, filter_width);
            }
        }

        let inv_total = 1.0f32 / conv.data.iter().sum::<f32>();
        for v in conv.data.iter_mut() {
            *v *= inv_total;
        }

        conv
    };

    // Compute the low-pass mask
    let mut low_pass_mask = Mask::new();
    for y in 0..MASK_SIZE {
        for x in 0..MASK_SIZE {
            for j in (-window_radius as isize)..(window_radius + 1) {
                for i in (-window_radius as isize)..(window_radius + 1) {
                    let b = y as isize + j;
                    let a = x as isize + i;
                    let alpha = conv[(i + window_radius, j + window_radius)];
                    low_pass_mask[(x as isize, y as isize)] += mask.get_wrapped(a, b) * alpha;
                }
            }
        }
    }

    // Subtract low pass from original
    for i in 0..mask.data.len() {
        mask.data[i] -= low_pass_mask.data[i];
    }
}

/// Remaps the values in a Mask to be linearly distributed within [0, 1]
fn remap_values(mask: &mut Mask) {
    let mut vals = Vec::with_capacity(MASK_SIZE * MASK_SIZE);
    for y in 0..MASK_SIZE {
        for x in 0..MASK_SIZE {
            vals.push((mask[(x as isize, y as isize)], x, y));
        }
    }
    vals.sort_by(|a, b| if a < b {
        Ordering::Less
    } else {
        Ordering::Greater
    });
    let inc = 1.0 / (vals.len() - 1) as f32;
    let mut n = 0.0;
    for v in vals.iter() {
        mask[(v.1 as isize, v.2 as isize)] = n;
        n += inc;
    }
}

/// Normal distribution
fn gauss(x: f32, sd: f32) -> f32 {
    let norm = 1.0 / (2.0 * std::f32::consts::PI * sd * sd).sqrt();
    let dist = ((-x * x) / (2.0 * sd * sd)).exp();
    norm * dist
}

// Sinc filter function
fn sinc(x: f32, w: f32) -> f32 {
    if x == 0.0 {
        1.0
    } else {
        let x = x * std::f32::consts::PI / w;
        x.sin() / x
    }
}

/// Returns a random float in [0, 1] based on 'n' and a seed.
/// Generally use n for getting a bunch of different random
/// numbers, and use seed to vary between runs.
pub fn hash_u32_to_f32(n: u32, seed: u32) -> f32 {
    let mut hash = n;

    for _ in 0..8 {
        hash = hash.wrapping_mul(1936502639);
        hash ^= hash.wrapping_shr(16);
        hash = hash.wrapping_add(seed);
    }

    const INV_MAX: f32 = 1.0 / std::u32::MAX as f32;
    return hash as f32 * INV_MAX;
}


// Holds data for a 2d mask
struct Mask {
    data: Vec<f32>,
}

impl Mask {
    fn new() -> Mask {
        Mask { data: vec![0.0; MASK_SIZE * MASK_SIZE] }
    }

    fn get_wrapped(&self, ix: isize, iy: isize) -> f32 {
        let x = (ix + MASK_SIZE as isize) as usize & MASK_SIZE_BITMASK;
        let y = (iy + MASK_SIZE as isize) as usize & MASK_SIZE_BITMASK;

        self.data[(y << MASK_SIZE_POW) + x]
    }
}

impl Index<(isize, isize)> for Mask {
    type Output = f32;
    fn index(&self, index: (isize, isize)) -> &f32 {
        &self.data[index.1 as usize * MASK_SIZE + index.0 as usize]
    }
}

impl IndexMut<(isize, isize)> for Mask {
    fn index_mut(&mut self, index: (isize, isize)) -> &mut f32 {
        &mut self.data[index.1 as usize * MASK_SIZE + index.0 as usize]
    }
}