use std::sync::Arc;

use arrow::{
    array::{
        Array, ArrayRef, BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array,
        Int64Array, StringArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array,
    },
    datatypes::DataType,
};
use luminal::prelude::*;
use mpera_frontend::op::ReduceKind;

use crate::{
    Result,
    luminal_interop::arrays::arrow_array_to_f32,
    typecheck::{numeric_values_as_f64, reduce_output_dtype_for_array},
};

pub fn reduce_array(kind: &ReduceKind, array: &dyn Array) -> Result<ArrayRef> {
    match kind {
        ReduceKind::Count => Ok(Arc::new(UInt64Array::from(vec![Some(
            (array.len() - array.null_count()) as u64,
        )])) as ArrayRef),
        ReduceKind::Last => last_array(array),
        ReduceKind::Percentile { x } => percentile_array(array, *x),
        ReduceKind::Sum
        | ReduceKind::Product
        | ReduceKind::Mean
        | ReduceKind::Stdev
        | ReduceKind::Min
        | ReduceKind::Max => luminal_numeric_reduce_array(kind, array),
    }
}

fn luminal_numeric_reduce_array(kind: &ReduceKind, array: &dyn Array) -> Result<ArrayRef> {
    let out_dtype = reduce_output_dtype_for_array(kind, array)?;
    let values = arrow_array_to_f32(array)?.compact_valid_values();
    let value = luminal_reduce(kind, &values)?;

    macro_rules! cast_scalar {
        ($arr:ty, $cast:expr) => {
            Ok(Arc::new(<$arr>::from(vec![value.map($cast)])) as ArrayRef)
        };
    }

    match out_dtype {
        DataType::Float64 => cast_scalar!(Float64Array, f64::from),
        DataType::Float32 => Ok(Arc::new(Float32Array::from(vec![value])) as ArrayRef),
        DataType::Int16 => cast_scalar!(Int16Array, |v| v as i16),
        DataType::Int32 => cast_scalar!(Int32Array, |v| v as i32),
        DataType::Int64 => cast_scalar!(Int64Array, |v| v as i64),
        DataType::UInt8 => cast_scalar!(UInt8Array, |v| v as u8),
        DataType::UInt16 => cast_scalar!(UInt16Array, |v| v as u16),
        DataType::UInt32 => cast_scalar!(UInt32Array, |v| v as u32),
        DataType::UInt64 => cast_scalar!(UInt64Array, |v| v as u64),
        _ => Err(autark_error::Error::Reader(format!(
            "reduce {:?} does not support dtype {} yet",
            kind, out_dtype
        ))),
    }
}

fn luminal_reduce(kind: &ReduceKind, values: &[f32]) -> Result<Option<f32>> {
    if values.is_empty() {
        return Ok(None);
    }

    let mut cx = Graph::new();
    let input = cx.tensor(values.len());
    let out = match kind {
        ReduceKind::Sum => input.sum(0).output(),
        ReduceKind::Product => input.prod(0).output(),
        ReduceKind::Mean => input.mean(0).output(),
        ReduceKind::Stdev => input.std_options(0, 0).output(),
        ReduceKind::Min => input.min(0).output(),
        ReduceKind::Max => input.max(0).output(),
        _ => {
            return Err(autark_error::Error::Reader(format!(
                "reduce {:?} is not implemented yet",
                kind
            )));
        }
    };

    cx.build_search_space::<NativeRuntime>();
    let mut rt = cx.search(NativeRuntime::default(), 1);
    rt.set_data(input, values.to_vec());
    rt.execute(&cx.dyn_map);
    Ok(rt.get_f32(out.id).first().copied())
}

fn percentile_array(array: &dyn Array, x: f64) -> Result<ArrayRef> {
    let mut values = numeric_values_as_f64(array)?;
    if values.is_empty() {
        return Ok(Arc::new(Float64Array::from(vec![None])) as ArrayRef);
    }
    values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let rank = ((values.len() - 1) as f64 * x.clamp(0.0, 1.0)).round() as usize;
    Ok(Arc::new(Float64Array::from(vec![values.get(rank).copied()])) as ArrayRef)
}

fn last_array(array: &dyn Array) -> Result<ArrayRef> {
    macro_rules! last_primitive {
        ($arr:ty, $out:ty) => {{
            let array = array
                .as_any()
                .downcast_ref::<$arr>()
                .ok_or(autark_error::Error::UnsupportedArrowDataType)?;
            Ok(Arc::new(<$out>::from(vec![array.iter().flatten().last()])) as ArrayRef)
        }};
    }

    match array.data_type() {
        DataType::Boolean => {
            let array = array
                .as_any()
                .downcast_ref::<BooleanArray>()
                .ok_or(autark_error::Error::UnsupportedArrowDataType)?;
            Ok(Arc::new(BooleanArray::from(vec![array.iter().flatten().last()])) as ArrayRef)
        }
        DataType::Utf8 => {
            let array = array
                .as_any()
                .downcast_ref::<StringArray>()
                .ok_or(autark_error::Error::UnsupportedArrowDataType)?;
            Ok(Arc::new(StringArray::from(vec![array.iter().flatten().last()])) as ArrayRef)
        }
        DataType::Int16 => last_primitive!(Int16Array, Int16Array),
        DataType::Int32 => last_primitive!(Int32Array, Int32Array),
        DataType::Int64 => last_primitive!(Int64Array, Int64Array),
        DataType::UInt8 => last_primitive!(UInt8Array, UInt8Array),
        DataType::UInt16 => last_primitive!(UInt16Array, UInt16Array),
        DataType::UInt32 => last_primitive!(UInt32Array, UInt32Array),
        DataType::UInt64 => last_primitive!(UInt64Array, UInt64Array),
        DataType::Float32 => last_primitive!(Float32Array, Float32Array),
        DataType::Float64 => last_primitive!(Float64Array, Float64Array),
        _ => Err(autark_error::Error::UnsupportedArrowDataType),
    }
}
