use crate::groups::GroupContainer;
use crate::groups::GroupContainerNamedTuple;
use crate::math::{chunk_and_fortranize, dense_fold_change};
use crate::ranking::{rank_sum_and_ties, sort_along_axis_0_inplace};
use crate::sparse::types::SparseFloat;
use crate::stats::compute_pvalue;
use ndarray::ArrayViewMut1;
use ndarray::{Array1, Array2, ArrayView2, s};
use numpy::{PyArray1, PyArray2, PyReadonlyArray2};
use pyo3::PyAny;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::{Bound, PyResult, Python, pyfunction};

pub fn dense_ovo_kernel<D: SparseFloat>(
    sorted_controls: ArrayView2<D>,
    sorted_tgt: ArrayView2<D>,
    use_continuity: bool,
    tie_correct: bool,
    alternative: &String,
    mut p_values: ArrayViewMut1<f64>,
    mut u_stats: ArrayViewMut1<f64>,
) -> Result<(), String> {
    let n_ctrl = sorted_controls.dim().0 as f64;
    let n_cols = sorted_controls.dim().1 as f64;
    let n_tgt = sorted_tgt.dim().0 as f64;

    // declare placeholder for results
    // let mut u_stats = Array1::zeros(n_cols as usize);
    // let mut p_values = Array1::zeros(n_cols as usize);

    let n = n_ctrl + n_tgt;
    let mu = n_ctrl * n_tgt / 2.;

    let u_base = n_ctrl as f64 * n_tgt + n_tgt * (n_tgt + 1.) / 2.;
    for j in 0..n_cols as usize {
        let (rs, ts) = rank_sum_and_ties(sorted_controls.column(j), sorted_tgt.column(j));

        let u = u_base - rs;

        let contin_corr = if use_continuity { 0.5 } else { 0. };
        let pv = compute_pvalue(
            n_ctrl,
            n_tgt,
            n,
            if tie_correct { ts } else { 0. },
            u,
            mu,
            contin_corr,
            alternative,
        )?;
        p_values[j] = pv;
        u_stats[j] = u;
    }

    return Ok(());
}

#[pyfunction]
pub fn dense_ovo_kernel_rust<'py>(
    py: Python<'py>,
    sorted_controls: PyReadonlyArray2<'py, f32>,
    sorted_tgt: PyReadonlyArray2<'py, f32>,
    use_continuity: bool,
    tie_correct: bool,
    alternative: String,
) -> PyResult<(Bound<'py, PyArray1<f64>>, Bound<'py, PyArray1<f64>>)> {
    let sorted_controls = sorted_controls.as_array();
    let mut p_values = Array1::zeros(sorted_controls.dim().1);
    let mut u_stats = Array1::zeros(sorted_controls.dim().1);
    _ = dense_ovo_kernel(
        sorted_controls,
        sorted_tgt.as_array(),
        use_continuity,
        tie_correct,
        &alternative,
        p_values.view_mut(),
        u_stats.view_mut(),
    )
    .map_err(PyValueError::new_err)?;
    return Ok((
        PyArray1::from_array(py, &p_values),
        PyArray1::from_array(py, &u_stats),
    ));
}

pub fn dense_ovo_over_contiguous_col_chunk<D: SparseFloat>(
    x: ArrayView2<D>,
    chunk_lb: usize,
    chunk_ub: usize,
    grpc: GroupContainer,
    is_log1p: bool,
    use_continuity: bool,
    tie_correct: bool,
    alternative: &String,
) -> Result<(Array2<f64>, Array2<f64>, Array2<f32>), String> {
    if chunk_lb >= chunk_ub {
        return Err(format!(
            "Chunking error: lower bound ({}) is not smaller than upper bound ({}).",
            { chunk_lb },
            { chunk_ub }
        ));
    }

    if grpc.encoded_ref_group < 0 {
        return Err(format!(
            "Encoded ref group can not be negative. Received {}.",
            grpc.encoded_ref_group
        ));
    }
    let encoded_ref_group = grpc.encoded_ref_group as usize;

    // Chunk control cells out and sort in-place
    let ctrl_indices = grpc.indices.slice(s![
        grpc.indptr[encoded_ref_group]..grpc.indptr[encoded_ref_group + 1]
    ]);
    let mut ctrl_chunk = chunk_and_fortranize(&x, chunk_lb, chunk_ub, Some(ctrl_indices))?;
    sort_along_axis_0_inplace(ctrl_chunk.view_mut())?;

    // Initialize result placeholders
    let n_groups = grpc.counts.len();
    let n_cols = chunk_ub - chunk_lb;
    let mut p_values = Array2::<f64>::zeros((n_groups, n_cols));
    let mut u_stats = Array2::<f64>::zeros((n_groups, n_cols));

    for i in 0..n_groups {
        if i as isize == grpc.encoded_ref_group {
            p_values.row_mut(i).fill(1.);
            u_stats.row_mut(i).fill(-1.);
        } else {
            // Grab indices of the target group's cells
            let tgt_indices = grpc.indices.slice(s![grpc.indptr[i]..grpc.indptr[i + 1]]);
            // Chunk the target group's cells
            let mut tgt_chunk = chunk_and_fortranize(&x, chunk_lb, chunk_ub, Some(tgt_indices))?;
            // Sort them
            sort_along_axis_0_inplace(tgt_chunk.view_mut())?;

            // Compute p values and u-stats
            _ = dense_ovo_kernel(
                ctrl_chunk.view(),
                tgt_chunk.view(),
                use_continuity,
                tie_correct,
                alternative,
                p_values.row_mut(i),
                u_stats.row_mut(i),
            )?;

            // Fill in placeholders
            // p_values.row_mut(i).assign(&p);
            // u_stats.row_mut(i).assign(&u);
        }
    }
    let fc = dense_fold_change(x.slice(s![.., chunk_lb..chunk_ub]), &grpc, is_log1p, false)?;
    return Ok((p_values, u_stats, fc));
}

macro_rules! run_ovo_branch {
    ($py:expr, $x:expr, $chunk_lb:expr, $chunk_ub:expr, $grpc:expr, $is_log1p:expr, $use_continuity:expr, $tie_correct:expr, $alternative:expr, $dt:ty) => {{
        let x_pyarray = $x.extract::<PyReadonlyArray2<'py, $dt>>()?;
        let x = x_pyarray.as_array();
        $py.detach(|| {
            dense_ovo_over_contiguous_col_chunk(
                x.view(),
                $chunk_lb,
                $chunk_ub,
                $grpc,
                $is_log1p,
                $use_continuity,
                $tie_correct,
                &$alternative,
            )
        })
        .map_err(PyValueError::new_err)
    }};
}

type PyArr2<'py> = Bound<'py, PyArray2<f64>>;

#[pyfunction]
pub fn dense_ovo_over_contiguous_col_chunk_rust<'py>(
    py: Python<'py>,
    // x: PyReadonlyArray2<'py, f32>,
    x: Bound<'py, PyAny>,
    chunk_lb: usize,
    chunk_ub: usize,
    grpc: GroupContainerNamedTuple<'py>,
    is_log1p: bool,
    use_continuity: bool,
    tie_correct: bool,
    alternative: String,
) -> PyResult<(PyArr2<'py>, PyArr2<'py>, Bound<'py, PyArray2<f32>>)> {
    let grpc = grpc.as_group_container();
    // let x = x.as_array();
    let data_dtype: String = x.getattr("dtype")?.getattr("str")?.extract()?;
    let (p_values, u_stats, fc) = match data_dtype.as_str() {
        "f32" | "<f4" => run_ovo_branch!(
            py,
            x,
            chunk_lb,
            chunk_ub,
            grpc,
            is_log1p,
            use_continuity,
            tie_correct,
            alternative,
            f32
        ),
        "f64" | "<f8" => run_ovo_branch!(
            py,
            x,
            chunk_lb,
            chunk_ub,
            grpc,
            is_log1p,
            use_continuity,
            tie_correct,
            alternative,
            f64
        ),
        _ => Err(PyValueError::new_err(format!(
            "Input data should be f32 or f64, received {}",
            data_dtype
        ))),
    }?;
    return Ok((
        PyArray2::from_array(py, &p_values),
        PyArray2::from_array(py, &u_stats),
        PyArray2::from_array(py, &fc),
    ));
}
