#![feature(test)]

use pyo3::prelude::*;
use rayon::prelude::*;
use rand::prelude::*;
use pyo3::types::PyFunction;
use pyo3::types::PyTuple;
use rand;


trait Opti {
    fn min(
        self,
        f: impl Fn(f64, f64) -> f64 + std::marker::Sync,
        x_min: f64,
        x_max: f64,
        y_min: f64,
        y_max: f64,
    ) -> (f64, f64, f64);

}


/// n: 粒子个数,
/// c_1,c_2均为加速常数，通常在区间[0,2]内取值
/// m:迭代次数
#[derive(Clone, Copy)]
pub struct LiZiQunSuanFa {
    n: usize,
    c_1: f64,
    c_2: f64,
    m: usize,
}


impl Opti for LiZiQunSuanFa {
    fn min(
        self,
        f: impl Fn(f64, f64) -> f64 + std::marker::Sync,
        x_min: f64,
        x_max: f64,
        y_min: f64,
        y_max: f64,
    ) -> (f64, f64, f64) {
        use rand::thread_rng;
        use rand::distributions::Uniform;
        use rand::distributions::Distribution;
        // 随机器
        let mut rng = thread_rng();
        let between_x = Uniform::from(x_min..x_max);
        let between_y = Uniform::from(y_min..y_max);


        let r_1 = Uniform::from(0.0..self.c_1);
        let r_2 = Uniform::from(0.0..self.c_2);

        // 最大速度
        let v_x_max = 0.5 * (x_max - x_min);
        let v_y_max = 0.5 * (y_max - y_min);

        let mut f_min = f64::MAX;
        let mut x_f_min: f64 = Default::default();
        let mut y_f_min: f64 = Default::default();

        // 状态(x,y,v_x,v_y,min,x_min,y_min)
        let mut state: Vec<(f64, f64, f64, f64, f64, f64, f64)> = std::iter::repeat_with(|| {
            let x = between_x.sample(&mut rng);
            let y = between_y.sample(&mut rng);
            let f_x_y = f(x, y);
            if f_x_y < f_min {
                f_min = f_x_y;
                x_f_min = x;
                y_f_min = y;
            }
            (x, y, 0.0, 0.0, f(x, y), x, y)
        })
        .take(self.n)
        .collect();

        for _ in 0..self.m {
            state = (&mut state)
                .into_par_iter()
                .map(
                    |(mut x, mut y, mut v_x, mut v_y, mut min, mut x_min, mut y_min)| {
                        let mut rng = thread_rng();
                        let v_x_new = v_x
                            + r_1.sample(&mut rng) * (x_min - x)
                            + self.c_2 * r_2.sample(&mut rng) * (x_f_min - x);

                        if v_x_new.abs() > v_x_max {
                            v_x = v_x_max
                        } else {
                            v_x = v_x_new
                        }

                        let v_y_new = v_y
                            + r_1.sample(&mut rng) * (y_min - y)
                            + self.c_2 * r_2.sample(&mut rng);

                        if v_y_new.abs() > v_y_max {
                            v_y = v_y_max
                        } else {
                            v_y = v_y_new
                        }

                        let x_new = x + v_x;

                        if x_new < x_min {
                            x = x_min
                        } else if x_new > x_max {
                            x = x_max
                        } else {
                            x = x_new
                        }

                        let y_new = y + v_y;

                        if y_new < y_min {
                            y = y_min
                        } else if y_new > y_max {
                            y = y_max
                        } else {
                            y = y_new
                        }

                        let f_x_y = f(x, y);
                        if f_x_y < min {
                            min = f_x_y;
                            x_min = x;
                            y_min = y;
                        }

                        (x, y, v_x, v_y, min, x_min, y_min)
                    },
                )
                .collect();

            (f_min, x_f_min, y_f_min) = state.iter().fold((f_min, x_f_min, y_f_min), |acc, x| {
                if acc.0 < x.4 {
                    acc
                } else {
                    (x.4, x.5, x.6)
                }
            });
        }

        (x_f_min, y_f_min, f_min)
    }
}



/// Formats the sum of two numbers as string.
#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}


/// A Python module implemented in Rust.
#[pymodule]
fn opti_solve(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    extern crate test;
    use test::Bencher;

    fn f<const N: usize> (f:impl Fn([f64;N]) -> f64 + std::marker::Sync ){}


    #[test]
    fn test() {

        // #[bench]
        // fn bencher(b: &mut Bencher) {
        //     use std::mem::size_of_val;
        //     struct Zero;
        //     struct One;
        //     static ZERO: Zero = Zero;
        //     static ONE: One = One;
        //
        //     let (zero, one) = unsafe { (&ZERO as *const Zero as usize, &ONE as *const One as usize) };
        //
        //     println!(
        //         "{}\n{}\n{}\n{}\n{}",
        //         zero,
        //         one,
        //         one as isize - zero as isize,
        //         size_of_val(&ZERO),
        //         size_of_val(&ONE)
        //     );
        //
        //     let v = vec![7u128; 100000];
        //     b.iter(|| v.iter().sum::<u128>());
        // }
    }
}