use autark::mpera::uprogram as rust_uprogram;
use pyo3::prelude::*;
use pyo3::types::PyModule;

use crate::error::py_err;
use crate::program::PyProgram;

#[pyfunction]
fn dist_normalized(program: PyRef<'_, PyProgram>) -> PyResult<PyProgram> {
    let out = rust_uprogram::dist_normalized(program.inner_clone()).map_err(py_err)?;
    Ok(PyProgram::new(program.graph_id(), out))
}

#[pyfunction]
fn zscore(program: PyRef<'_, PyProgram>) -> PyResult<PyProgram> {
    let out = rust_uprogram::zscore(program.inner_clone()).map_err(py_err)?;
    Ok(PyProgram::new(program.graph_id(), out))
}

pub(crate) fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> {
    let py = parent.py();
    let sub = PyModule::new(py, "uprogram")?;
    sub.add_function(wrap_pyfunction!(dist_normalized, &sub)?)?;
    sub.add_function(wrap_pyfunction!(zscore, &sub)?)?;
    parent.add_submodule(&sub)?;

    // Make `from autark.uprogram import *` work for this extension module.
    let sys = py.import("sys")?;
    let modules = sys.getattr("modules")?;
    modules.set_item("autark.uprogram", &sub)?;
    Ok(())
}
