extern crate time;
extern crate simple_parallel;
extern crate parry;

use parry::{Expression, E, Constant};
use time::Duration;
use std::cmp;

fn naive(a: &[f64], b: &[f64]) -> Vec<f64> {
    let na = a.len();
    let nb = b.len();
    let nc = na + nb - 1;

    let mut v = vec![0.0; nc];
    let mut temporary = vec![0.0; cmp::min(na, nb)];

    for (k, place) in v.iter_mut().enumerate() {
        let lo_j = if k + 1 > nb { k - nb + 1 } else { 0 };
        let hi_j = cmp::min(k + 1, na);

        let slice_a = &a[lo_j..hi_j];
        let slice_b = &b[(k + 1) - hi_j..k - lo_j +1];
        for (o, (aa, bb)) in temporary.iter_mut().zip(slice_a.iter().zip(slice_b.iter().rev())) {
            *o = aa + bb;
        }
        let tlen = hi_j - lo_j;
        let tmp = &mut temporary[..tlen];

        let mut max = std::f64::NEG_INFINITY;
        let mut max_i = 0;
        for (i, t) in tmp.iter().enumerate() {
            if *t > max {
                max = *t;
                max_i = i;
            }
        }
        if max == std::f64::NEG_INFINITY {
            *place = max;
            continue
        }
        for t in &mut *tmp {
            *t = t.exp() - max
        }

        *place = (tmp[..max_i].iter().fold(0.0, |a, &x| a + x) +
                  tmp[max_i + 1..].iter().fold(0.0, |a, &x| a + x)).ln_1p() + max;
    }
    v
}

fn sp(a: &[f64], b: &[f64]) -> Vec<f64> {
    let na = a.len();
    let nb = b.len();
    let nc = na + nb - 1;

    let mut v = vec![0.0; nc];

    simple_parallel::Pool::new(8).for_(v.iter_mut().enumerate(), |(k, place)| {
        let lo_j = if k + 1 > nb { k - nb + 1 } else { 0 };
        let hi_j = cmp::min(k + 1, na);

        let slice_a = &a[lo_j..hi_j];
        let slice_b = &b[(k + 1) - hi_j..k - lo_j +1];
        let mut tmp = vec![0.0; hi_j - lo_j];
        for (o, (aa, bb)) in tmp.iter_mut().zip(slice_a.iter().zip(slice_b.iter().rev())) {
            *o = aa + bb;
        }
        let mut max = std::f64::NEG_INFINITY;
        let mut max_i = 0;
        for (i, t) in tmp.iter().enumerate() {
            if *t > max {
                max = *t;
                max_i = i;
            }
        }
        if max == std::f64::NEG_INFINITY {
            *place = max;
            return
        }
        for t in &mut tmp {
            *t = (*t - max).exp()
        }

        *place = (tmp[..max_i].iter().fold(0.0, |a, &x| a + x) +
                  tmp[max_i + 1..].iter().fold(0.0, |a, &x| a + x)).ln_1p() + max;
    });

    v
}

fn p(a: &[f64], b: &[f64]) -> Vec<f64> {
    let na = a.len();
    let nb = b.len();
    let nc = na + nb - 1;

    let mut v = vec![0.0; nc];

    let mut temporary1 = vec![0.0; cmp::min(na, nb)];
    let mut temporary2 = vec![0.0; cmp::min(na, nb)];

    for (k, place) in v.iter_mut().enumerate() {
        let lo_j = if k + 1 > nb { k - nb + 1 } else { 0 };
        let hi_j = cmp::min(k + 1, na);

        let slice_a = &a[lo_j..hi_j];
        let slice_b = &b[(k + 1) - hi_j..k - lo_j +1];
        let tlen = hi_j - lo_j;
        let tmp = &mut temporary1[..tlen];
        let tmp2 = &mut temporary2[..tlen];
        parry::evaluate(tmp, slice_b.rev() + slice_a);
        let max = E(&*tmp).max().unwrap_or(std::f64::NEG_INFINITY);
        if max == std::f64::NEG_INFINITY {
            *place = max;
            continue
        }
        parry::evaluate(tmp2, (E(&*tmp) - Constant(max)).map(&f64::exp));

        *place = E(&*tmp2).sum() + max;
    }
    v
}

fn measure<T, F: FnMut() -> T>(name: &str, mut f: F) {
    let time = Duration::span(|| {
        f();
    });

    let secs = time.num_nanoseconds().unwrap() as f64 * 1e-9;
    println!("{}: {:.3}s", name, secs);
}

fn main() {
    let vec = vec![0.0; 4000];

    measure("naive", || naive(&vec, &vec));
    measure("simple_parallel", || sp(&vec, &vec));
    measure("parry", || p(&vec, &vec));
}