use numpy::{PyArray1, get_array_module};
use pyo3::prelude::{ObjectProtocol, PyResult, Python};
use pyo3::types::PyDict;

fn main() -> Result<(), ()> {
    let gil = Python::acquire_gil();
    main_(gil.python()).map_err(|e| {
        eprintln!("error! :{:?}", e);
        // we can't display python error type via ::std::fmt::Display
        // so print error here manually
        e.print_and_set_sys_last_vars(gil.python());
    })
}

fn main_<'py>(py: Python<'py>) -> PyResult<()> {
    let np = py.import("numpy")?;
    let dict = PyDict::new(py);
    dict.set_item("np", np)?;
    let pyarray: &PyArray1<i32> = py
        .eval("np.absolute(np.array([-1, -2, -3], dtype='int32'))", Some(&dict), None)?
        .extract()?;
    let slice = pyarray.as_slice()?;
    assert_eq!(slice, &[1, 2, 3]);
    Ok(())
}

use ndarray::{ArrayD, ArrayViewD, ArrayViewMutD};
use numpy::{IntoPyArray, PyArrayDyn};
use pyo3::prelude::{pymodule, Py, PyModule, PyResult, Python};

#[pymodule]
fn rust_ext(_py: Python, m: &PyModule) -> PyResult<()> {
    // immutable example
    fn axpy(a: f64, x: ArrayViewD<f64>, y: ArrayViewD<f64>) -> ArrayD<f64> {
        a * &x + &y
    }

    // mutable example (no return)
    fn mult(a: f64, mut x: ArrayViewMutD<f64>) {
        x *= a;
    }

    // wrapper of `axpy`
    #[pyfn(m, "axpy")]
    fn axpy_py(
        py: Python,
        a: f64,
        x: &PyArrayDyn<f64>,
        y: &PyArrayDyn<f64>,
    ) -> Py<PyArrayDyn<f64>> {
        let x = x.as_array();
        let y = y.as_array();
        axpy(a, x, y).into_pyarray(py).to_owned()
    }

    // wrapper of `mult`
    #[pyfn(m, "mult")]
    fn mult_py(_py: Python, a: f64, x: &PyArrayDyn<f64>) -> PyResult<()> {
        let x = x.as_array_mut();
        mult(a, x);
        Ok(())
    }

    Ok(())
}