use std::sync::Arc;

use arrow::{
    array::{Array, ArrayRef, Float32Array, Float64Array},
    datatypes::DataType,
};
use luminal::prelude::*;
use mpera_frontend::op::UnaryOpKind;

use crate::{
    Result,
    luminal_interop::arrays::{MaskedF32, arrow_array_to_f32},
};

pub fn unary_array(kind: &UnaryOpKind, array: &dyn Array) -> Result<ArrayRef> {
    let out_dtype = match kind {
        UnaryOpKind::Neg | UnaryOpKind::Cumsum => array.data_type().clone(),
        _ => {
            return Err(autark_error::Error::Reader(format!(
                "unary {:?} does not use luminal unary interop",
                kind
            )));
        }
    };

    let input = arrow_array_to_f32(array)?;
    let values = luminal_unary(kind, &input.compact_valid_values())?;

    macro_rules! cast_with_validity {
        ($arr:ty, $cast:expr) => {
            Ok(Arc::new(<$arr>::from(
                with_validity(&input, values)
                    .into_iter()
                    .map(|value| value.map($cast))
                    .collect::<Vec<_>>(),
            )) as ArrayRef)
        };
    }

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

fn luminal_unary(kind: &UnaryOpKind, values: &[f32]) -> Result<Vec<f32>> {
    let mut cx = Graph::new();
    let input = cx.tensor(values.len());
    let out = match kind {
        UnaryOpKind::Neg => (-input).output(),
        UnaryOpKind::Cumsum => input.cumsum(0).output(),
        _ => {
            return Err(autark_error::Error::Reader(format!(
                "unary {:?} is not implemented in luminal interop",
                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).clone())
}

fn with_validity(input: &MaskedF32, values: Vec<f32>) -> Vec<Option<f32>> {
    let mut valid_values = values.into_iter();
    input
        .validity
        .iter()
        .map(|valid| valid.then(|| valid_values.next().unwrap()))
        .collect()
}
