//! MATLAB-compatible `range` builtin with GPU-aware semantics for RunMat.

use std::collections::HashSet;

use runmat_accelerate_api::GpuTensorHandle;
use runmat_builtins::{ResolveContext, Tensor, Type, Value};
use runmat_macros::runtime_builtin;

use crate::build_runtime_error;
use crate::builtins::common::arg_tokens::tokens_from_values;
use crate::builtins::common::spec::{
    BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
    ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
};
use crate::builtins::common::{gpu_helpers, tensor};

#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::array::creation::range")]
pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
    name: "range",
    op_kind: GpuOpKind::Reduction,
    supported_precisions: &[ScalarType::F32, ScalarType::F64],
    broadcast: BroadcastSemantics::Matlab,
    provider_hooks: &[
        ProviderHook::Reduction {
            name: "reduce_min_dim",
        },
        ProviderHook::Reduction {
            name: "reduce_max_dim",
        },
        ProviderHook::Reduction {
            name: "reduce_min",
        },
        ProviderHook::Reduction {
            name: "reduce_max",
        },
        ProviderHook::Binary {
            name: "elem_sub",
            commutative: false,
        },
    ],
    constant_strategy: ConstantStrategy::InlineLiteral,
    residency: ResidencyPolicy::NewHandle,
    nan_mode: ReductionNaN::Include,
    two_pass_threshold: Some(256),
    workgroup_size: Some(256),
    accepts_nan_mode: false,
    notes: "Requires provider min/max reductions plus elem_sub; omitnan and multi-axis reductions gather to host when hooks are absent.",
};

fn builtin_error(message: impl Into<String>) -> crate::RuntimeError {
    build_runtime_error(message).with_builtin("range").build()
}

fn range_type(args: &[Type], _context: &ResolveContext) -> Type {
    let Some(input) = args.first() else {
        return Type::Unknown;
    };
    match input {
        Type::Num | Type::Int | Type::Bool => Type::Num,
        Type::Tensor { .. } | Type::Logical { .. } => Type::tensor(),
        _ => Type::tensor(),
    }
}

#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::array::creation::range")]
pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
    name: "range",
    shape: ShapeRequirements::BroadcastCompatible,
    constant_strategy: ConstantStrategy::InlineLiteral,
    elementwise: None,
    reduction: None,
    emits_nan: true,
    notes: "No fused kernel today; the runtime composes provider min/max reductions or gathers to host.",
};

#[runtime_builtin(
    name = "range",
    category = "array/creation",
    summary = "Compute the difference between the maximum and minimum values.",
    keywords = "range,max,min,spread,gpu",
    accel = "reduction",
    type_resolver(range_type),
    builtin_path = "crate::builtins::array::creation::range"
)]
async fn range_builtin(value: Value, rest: Vec<Value>) -> crate::BuiltinResult<Value> {
    let (dim_selection, nan_mode) = parse_arguments(&rest)?;
    match value {
        Value::GpuTensor(handle) => range_gpu(handle, dim_selection, nan_mode).await,
        other => range_host(other, dim_selection, nan_mode),
    }
}

#[derive(Clone, Copy, PartialEq, Eq)]
enum NanMode {
    Include,
    Omit,
}

#[derive(Clone, PartialEq, Eq)]
enum DimSelection {
    Auto,
    Dim(usize),
    Vec(Vec<usize>),
    All,
}

#[derive(Clone)]
struct ResolvedDims {
    dims_in_bounds: Vec<usize>,
    #[allow(dead_code)]
    dims_out_of_bounds: Vec<usize>,
}

fn parse_arguments(args: &[Value]) -> crate::BuiltinResult<(DimSelection, NanMode)> {
    let mut selection = DimSelection::Auto;
    let mut nan_mode = NanMode::Include;
    let mut selection_set = false;

    let tokens = tokens_from_values(args);
    for (arg, token) in args.iter().zip(tokens.iter()) {
        if let crate::builtins::common::arg_tokens::ArgToken::String(text) = token {
            match text.as_str() {
                "omitnan" => {
                    nan_mode = NanMode::Omit;
                    continue;
                }
                "includenan" => {
                    nan_mode = NanMode::Include;
                    continue;
                }
                "all" => {
                    if selection_set && !matches!(selection, DimSelection::Auto) {
                        return Err(builtin_error(
                            "range: 'all' cannot be combined with an explicit dimension",
                        ));
                    }
                    selection = DimSelection::All;
                    selection_set = true;
                    continue;
                }
                _ => {}
            }
        }

        if let Some(mode) = parse_nan_flag(arg)? {
            nan_mode = mode;
            continue;
        }

        if is_all_flag(arg)? {
            if selection_set && !matches!(selection, DimSelection::Auto) {
                return Err(builtin_error(
                    "range: 'all' cannot be combined with an explicit dimension",
                ));
            }
            selection = DimSelection::All;
            selection_set = true;
            continue;
        }

        if selection_set && !matches!(selection, DimSelection::Auto) {
            return Err(builtin_error("range: too many dimension arguments"));
        }

        selection = parse_dim_spec(arg)?;
        selection_set = true;
    }

    Ok((selection, nan_mode))
}

fn parse_nan_flag(value: &Value) -> crate::BuiltinResult<Option<NanMode>> {
    let text = match value {
        Value::String(s) => Some(s.clone()),
        Value::StringArray(sa) if sa.data.len() == 1 => Some(sa.data[0].clone()),
        Value::CharArray(ca) if ca.rows == 1 => Some(ca.data.iter().collect()),
        _ => None,
    };
    let Some(text) = text else {
        return Ok(None);
    };
    let lowered = text.trim().to_ascii_lowercase();
    match lowered.as_str() {
        "omitnan" => Ok(Some(NanMode::Omit)),
        "includenan" => Ok(Some(NanMode::Include)),
        _ => Ok(None),
    }
}

fn is_all_flag(value: &Value) -> crate::BuiltinResult<bool> {
    let text = match value {
        Value::String(s) => Some(s.clone()),
        Value::StringArray(sa) if sa.data.len() == 1 => Some(sa.data[0].clone()),
        Value::CharArray(ca) if ca.rows == 1 => Some(ca.data.iter().collect()),
        _ => None,
    };
    Ok(text
        .map(|s| s.trim().eq_ignore_ascii_case("all"))
        .unwrap_or(false))
}

fn parse_dim_spec(value: &Value) -> crate::BuiltinResult<DimSelection> {
    match value {
        Value::Int(i) => {
            let dim = i.to_i64();
            if dim < 1 {
                return Err(builtin_error("range: dimension must be >= 1"));
            }
            Ok(DimSelection::Dim(dim as usize))
        }
        Value::Num(n) => {
            if !n.is_finite() {
                return Err(builtin_error("range: dimension must be finite"));
            }
            let rounded = n.round();
            if (rounded - n).abs() > f64::EPSILON {
                return Err(builtin_error("range: dimension must be an integer"));
            }
            if rounded < 1.0 {
                return Err(builtin_error("range: dimension must be >= 1"));
            }
            Ok(DimSelection::Dim(rounded as usize))
        }
        Value::Tensor(t) => parse_dim_tensor(t),
        Value::LogicalArray(logical) => {
            let tensor = tensor::logical_to_tensor(logical)
                .map_err(|e| builtin_error(format!("range: {e}")))?;
            parse_dim_tensor(&tensor)
        }
        Value::GpuTensor(_) => Err(builtin_error(
            "range: dimension arguments must reside on the host (numeric or string)",
        )),
        other => Err(builtin_error(format!(
            "range: unsupported dimension argument type {:?}",
            other
        ))),
    }
}

fn parse_dim_tensor(tensor: &Tensor) -> crate::BuiltinResult<DimSelection> {
    if tensor.data.is_empty() {
        return Ok(DimSelection::Auto);
    }
    if !is_vector_shape(&tensor.shape) {
        return Err(builtin_error(
            "range: dimension vector must be a row or column vector",
        ));
    }
    let mut dims = Vec::with_capacity(tensor.data.len());
    for &value in &tensor.data {
        if !value.is_finite() {
            return Err(builtin_error("range: dimensions must be finite"));
        }
        let rounded = value.round();
        if (rounded - value).abs() > f64::EPSILON {
            return Err(builtin_error("range: dimensions must contain integers"));
        }
        if rounded < 1.0 {
            return Err(builtin_error("range: dimension indices must be >= 1"));
        }
        dims.push(rounded as usize);
    }
    Ok(DimSelection::Vec(dims))
}

fn is_vector_shape(shape: &[usize]) -> bool {
    match shape.len() {
        0 => true,
        1 => true,
        2 => shape[0] == 1 || shape[1] == 1,
        _ => shape.iter().filter(|&&d| d > 1).count() <= 1,
    }
}

fn range_host(
    value: Value,
    selection: DimSelection,
    nan_mode: NanMode,
) -> crate::BuiltinResult<Value> {
    let tensor = tensor::value_into_tensor_for("range", value)
        .map_err(|e| builtin_error(format!("range: {e}")))?;
    let resolved = resolve_dims(&tensor.shape, &selection)?;
    let result = compute_range_tensor(&tensor, &resolved, nan_mode)?;
    Ok(tensor::tensor_into_value(result))
}

async fn range_gpu(
    handle: GpuTensorHandle,
    selection: DimSelection,
    nan_mode: NanMode,
) -> crate::BuiltinResult<Value> {
    if matches!(nan_mode, NanMode::Omit) {
        return range_gpu_fallback(&handle, selection, nan_mode).await;
    }

    #[cfg(all(test, feature = "wgpu"))]
    {
        if handle.device_id != 0 {
            let _ = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
                runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
            );
        }
    }
    let Some(provider) = runmat_accelerate_api::provider() else {
        return range_gpu_fallback(&handle, selection, nan_mode).await;
    };

    let resolved = resolve_dims(&handle.shape, &selection)?;

    if should_use_global_reduce(&handle.shape, &selection, &resolved) {
        if let Some(diff) = range_gpu_all(provider, &handle).await {
            return Ok(Value::GpuTensor(diff));
        }
        return range_gpu_fallback(&handle, selection, nan_mode).await;
    }

    if resolved.dims_in_bounds.len() != 1 {
        return range_gpu_fallback(&handle, selection, nan_mode).await;
    }

    let dim = resolved.dims_in_bounds[0];
    let expected_shape = reduced_shape(&handle.shape, &[dim]);

    if expected_shape.is_empty() {
        return range_gpu_fallback(&handle, selection, nan_mode).await;
    }

    if let Some(diff) = range_gpu_single_dim(provider, &handle, dim, &expected_shape).await {
        return Ok(Value::GpuTensor(diff));
    }

    range_gpu_fallback(&handle, selection, nan_mode).await
}

fn reduced_shape(shape: &[usize], dims: &[usize]) -> Vec<usize> {
    if shape.is_empty() {
        return Vec::new();
    }
    let mut out = shape.to_vec();
    for &dim in dims {
        if dim < out.len() {
            out[dim] = 1;
        }
    }
    out
}

fn should_use_global_reduce(
    shape: &[usize],
    selection: &DimSelection,
    resolved: &ResolvedDims,
) -> bool {
    if shape.is_empty() {
        return false;
    }
    matches!(selection, DimSelection::All) || resolved.dims_in_bounds.len() == shape.len()
}

async fn range_gpu_all(
    provider: &'static dyn runmat_accelerate_api::AccelProvider,
    handle: &GpuTensorHandle,
) -> Option<GpuTensorHandle> {
    let min_handle = provider.reduce_min(handle).await.ok()?;
    let max_handle = match provider.reduce_max(handle).await {
        Ok(h) => h,
        Err(_) => {
            let _ = provider.free(&min_handle);
            return None;
        }
    };
    let diff = match provider.elem_sub(&max_handle, &min_handle).await {
        Ok(h) => h,
        Err(_) => {
            let _ = provider.free(&min_handle);
            let _ = provider.free(&max_handle);
            return None;
        }
    };
    let _ = provider.free(&min_handle);
    let _ = provider.free(&max_handle);
    Some(diff)
}

async fn range_gpu_single_dim(
    provider: &'static dyn runmat_accelerate_api::AccelProvider,
    handle: &GpuTensorHandle,
    dim_zero_based: usize,
    expected_shape: &[usize],
) -> Option<GpuTensorHandle> {
    let mut candidates = Vec::with_capacity(2);
    candidates.push(dim_zero_based);
    if let Some(next) = dim_zero_based.checked_add(1) {
        if next != dim_zero_based {
            candidates.push(next);
        }
    }
    let mut seen = HashSet::new();
    for candidate in candidates {
        if !seen.insert(candidate) {
            continue;
        }
        let min_result = match provider.reduce_min_dim(handle, candidate).await {
            Ok(res) => res,
            Err(_) => continue,
        };
        if min_result.values.shape != expected_shape {
            let _ = provider.free(&min_result.values);
            let _ = provider.free(&min_result.indices);
            continue;
        }
        let max_result = match provider.reduce_max_dim(handle, candidate).await {
            Ok(res) => res,
            Err(_) => {
                let _ = provider.free(&min_result.values);
                let _ = provider.free(&min_result.indices);
                continue;
            }
        };
        if max_result.values.shape != expected_shape {
            let _ = provider.free(&min_result.values);
            let _ = provider.free(&min_result.indices);
            let _ = provider.free(&max_result.values);
            let _ = provider.free(&max_result.indices);
            continue;
        }
        let diff = match provider
            .elem_sub(&max_result.values, &min_result.values)
            .await
        {
            Ok(handle) => handle,
            Err(_) => {
                let _ = provider.free(&min_result.values);
                let _ = provider.free(&min_result.indices);
                let _ = provider.free(&max_result.values);
                let _ = provider.free(&max_result.indices);
                continue;
            }
        };
        let _ = provider.free(&min_result.values);
        let _ = provider.free(&min_result.indices);
        let _ = provider.free(&max_result.values);
        let _ = provider.free(&max_result.indices);
        return Some(diff);
    }
    None
}

async fn range_gpu_fallback(
    handle: &GpuTensorHandle,
    selection: DimSelection,
    nan_mode: NanMode,
) -> crate::BuiltinResult<Value> {
    let tensor = gpu_helpers::gather_tensor_async(handle).await?;
    range_host(Value::Tensor(tensor), selection, nan_mode)
}

fn resolve_dims(shape: &[usize], selection: &DimSelection) -> crate::BuiltinResult<ResolvedDims> {
    let dims_1_based: Vec<usize> = match selection {
        DimSelection::Auto => vec![default_dimension_from_shape(shape)],
        DimSelection::Dim(d) => vec![*d],
        DimSelection::Vec(v) => {
            if v.is_empty() {
                vec![default_dimension_from_shape(shape)]
            } else {
                v.clone()
            }
        }
        DimSelection::All => {
            let ndims = if shape.is_empty() { 1 } else { shape.len() };
            (1..=ndims).collect()
        }
    };

    let mut seen = HashSet::new();
    let mut dims_in_bounds = Vec::new();
    let mut dims_out_of_bounds = Vec::new();
    let ndims = shape.len();

    for dim1 in dims_1_based {
        if dim1 == 0 {
            return Err(builtin_error("range: dimension indices must be >= 1"));
        }
        if !seen.insert(dim1) {
            continue;
        }
        let zero = dim1 - 1;
        if zero < ndims {
            dims_in_bounds.push(zero);
        } else {
            dims_out_of_bounds.push(dim1);
        }
    }

    dims_in_bounds.sort_unstable();
    dims_out_of_bounds.sort_unstable();

    Ok(ResolvedDims {
        dims_in_bounds,
        dims_out_of_bounds,
    })
}

fn compute_range_tensor(
    tensor: &Tensor,
    dims: &ResolvedDims,
    nan_mode: NanMode,
) -> crate::BuiltinResult<Tensor> {
    let mut shape = tensor.shape.clone();
    if shape.is_empty() {
        shape = vec![tensor.rows, tensor.cols];
    }

    if dims.dims_in_bounds.is_empty() {
        let output_shape = shape.clone();
        let mut output = Vec::with_capacity(tensor.data.len());
        for &value in &tensor.data {
            if value.is_nan() {
                output.push(f64::NAN);
            } else {
                output.push(0.0);
            }
        }
        return Tensor::new(output, output_shape).map_err(|e| builtin_error(format!("range: {e}")));
    }

    let mut output_shape = shape.clone();
    for &dim in &dims.dims_in_bounds {
        if dim < output_shape.len() {
            output_shape[dim] = 1;
        }
    }

    let out_len = tensor::element_count(&output_shape);
    if out_len == 0 {
        return Tensor::new(vec![], output_shape).map_err(|e| builtin_error(format!("range: {e}")));
    }

    let mut mins = vec![f64::INFINITY; out_len];
    let mut maxs = vec![f64::NEG_INFINITY; out_len];
    let mut saw_value = vec![false; out_len];
    let mut saw_nan = vec![false; out_len];

    let mut coords = vec![0usize; shape.len()];
    let mut out_coords = vec![0usize; shape.len()];
    let mut reduce_mask = vec![false; shape.len()];
    for &dim in &dims.dims_in_bounds {
        if dim < reduce_mask.len() {
            reduce_mask[dim] = true;
        }
    }

    for (linear, &value) in tensor.data.iter().enumerate() {
        linear_to_multi(linear, &shape, &mut coords);
        for (i, coord) in coords.iter().enumerate() {
            out_coords[i] = if reduce_mask[i] { 0 } else { *coord };
        }
        let out_idx = multi_to_linear(&out_coords, &output_shape);
        if value.is_nan() {
            if matches!(nan_mode, NanMode::Include) {
                saw_nan[out_idx] = true;
            }
            continue;
        }
        if !saw_value[out_idx] {
            mins[out_idx] = value;
            maxs[out_idx] = value;
            saw_value[out_idx] = true;
        } else {
            if value < mins[out_idx] {
                mins[out_idx] = value;
            }
            if value > maxs[out_idx] {
                maxs[out_idx] = value;
            }
        }
    }

    let mut output = vec![0.0; out_len];
    for idx in 0..out_len {
        if saw_nan[idx] {
            output[idx] = f64::NAN;
        } else if saw_value[idx] {
            let diff = maxs[idx] - mins[idx];
            output[idx] = if diff == 0.0 { 0.0 } else { diff };
        } else {
            output[idx] = f64::NAN;
        }
    }

    Tensor::new(output, output_shape).map_err(|e| builtin_error(format!("range: {e}")))
}

fn linear_to_multi(index: usize, shape: &[usize], out: &mut [usize]) {
    let mut remainder = index;
    for (dim, &size) in shape.iter().enumerate() {
        if size == 0 {
            out[dim] = 0;
        } else {
            out[dim] = remainder % size;
            remainder /= size;
        }
    }
}

fn multi_to_linear(coords: &[usize], shape: &[usize]) -> usize {
    let mut stride = 1usize;
    let mut index = 0usize;
    for (dim, &size) in shape.iter().enumerate() {
        if size == 0 {
            continue;
        }
        index += coords[dim] * stride;
        stride *= size;
    }
    index
}

fn default_dimension_from_shape(shape: &[usize]) -> usize {
    if shape.is_empty() {
        return 1;
    }
    shape
        .iter()
        .position(|&extent| extent != 1)
        .map(|idx| idx + 1)
        .unwrap_or(1)
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use crate::builtins::common::test_support;
    use futures::executor::block_on;

    fn range_builtin(value: Value, rest: Vec<Value>) -> crate::BuiltinResult<Value> {
        block_on(super::range_builtin(value, rest))
    }
    use runmat_accelerate_api::HostTensorView;
    use runmat_builtins::IntValue;

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_scalar_zero() {
        let result = range_builtin(Value::Num(42.0), Vec::new()).expect("range");
        match result {
            Value::Num(n) => assert_eq!(n, 0.0),
            other => panic!("expected scalar zero, got {other:?}"),
        }
    }

    #[test]
    fn range_type_scalar_returns_num() {
        assert_eq!(
            range_type(&[Type::Num], &ResolveContext::new(Vec::new())),
            Type::Num
        );
    }

    #[test]
    fn range_type_tensor_returns_tensor() {
        let out = range_type(
            &[Type::Tensor { shape: None }],
            &ResolveContext::new(Vec::new()),
        );
        assert_eq!(out, Type::tensor());
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_vector_default_dimension() {
        let tensor = Tensor::new(vec![1.0, 4.0, 2.0], vec![1, 3]).unwrap();
        let result = range_builtin(Value::Tensor(tensor), Vec::new()).expect("range");
        match result {
            Value::Num(n) => assert_eq!(n, 3.0),
            other => panic!("expected scalar result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_matrix_columnwise() {
        let tensor = Tensor::new(vec![1.0, 3.0, 4.0, 7.0, 2.0, 5.0], vec![2, 3]).unwrap();
        let result = range_builtin(Value::Tensor(tensor), Vec::new()).expect("range");
        match result {
            Value::Tensor(t) => {
                assert_eq!(t.shape, vec![1, 3]);
                assert_eq!(t.data, vec![2.0, 3.0, 3.0]);
            }
            other => panic!("expected row vector, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_matrix_rowwise() {
        let tensor = Tensor::new(vec![1.0, 3.0, 4.0, 7.0, 2.0, 5.0], vec![2, 3]).unwrap();
        let args = vec![Value::Int(IntValue::I32(2))];
        let result = range_builtin(Value::Tensor(tensor), args).expect("range");
        match result {
            Value::Tensor(t) => {
                assert_eq!(t.shape, vec![2, 1]);
                assert_eq!(t.data, vec![3.0, 4.0]);
            }
            other => panic!("expected column vector, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_all_collapse() {
        let tensor =
            Tensor::new((1..=6).map(|v| v as f64).collect::<Vec<_>>(), vec![2, 3]).unwrap();
        let args = vec![Value::from("all")];
        let result = range_builtin(Value::Tensor(tensor), args).expect("range");
        match result {
            Value::Num(n) => assert_eq!(n, 5.0),
            other => panic!("expected scalar result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_omit_nan() {
        let tensor = Tensor::new(vec![2.0, 4.0, f64::NAN, 6.0, 5.0, f64::NAN], vec![2, 3]).unwrap();
        let args = vec![Value::Int(IntValue::I32(1)), Value::from("omitnan")];
        let result = range_builtin(Value::Tensor(tensor), args).expect("range");
        match result {
            Value::Tensor(t) => {
                assert_eq!(t.shape, vec![1, 3]);
                assert_eq!(t.data, vec![2.0, 0.0, 0.0]);
            }
            other => panic!("expected row vector, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_includenan_returns_nan() {
        let tensor = Tensor::new(vec![2.0, f64::NAN, 5.0], vec![3, 1]).unwrap();
        let result =
            range_builtin(Value::Tensor(tensor), vec![Value::from("includenan")]).expect("range");
        match result {
            Value::Num(n) => assert!(n.is_nan()),
            other => panic!("expected scalar NaN, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_all_omitnan() {
        let tensor = Tensor::new(vec![2.0, f64::NAN, 10.0, 4.0], vec![2, 2]).unwrap();
        let args = vec![Value::from("all"), Value::from("omitnan")];
        let result = range_builtin(Value::Tensor(tensor), args).expect("range");
        match result {
            Value::Num(n) => assert_eq!(n, 8.0),
            other => panic!("expected scalar result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_dim_beyond_ndims_returns_zeros_or_nan() {
        let tensor = Tensor::new(vec![1.0, 4.0, f64::NAN, 7.0], vec![2, 2]).unwrap();
        let args = vec![Value::Int(IntValue::I32(3))];
        let result = range_builtin(Value::Tensor(tensor), args).expect("range");
        match result {
            Value::Tensor(t) => {
                assert_eq!(t.shape, vec![2, 2]);
                assert_eq!(t.data[0], 0.0);
                assert_eq!(t.data[1], 0.0);
                assert!(t.data[2].is_nan());
                assert_eq!(t.data[3], 0.0);
            }
            other => panic!("expected matrix result, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_multiple_dimensions() {
        let tensor = Tensor::new(
            (1..=24).map(|v| v as f64).collect::<Vec<_>>(),
            vec![3, 4, 2],
        )
        .unwrap();
        let dims = Tensor::new(vec![1.0, 2.0], vec![1, 2]).unwrap();
        let result =
            range_builtin(Value::Tensor(tensor), vec![Value::Tensor(dims)]).expect("range");
        match result {
            Value::Tensor(t) => {
                assert_eq!(t.shape, vec![1, 1, 2]);
                assert_eq!(t.data, vec![11.0, 11.0]);
            }
            other => panic!("expected 1x1x2 tensor, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_vecdim_column_vector() {
        let tensor = Tensor::new(
            (1..=12).map(|v| v as f64).collect::<Vec<_>>(),
            vec![3, 2, 2],
        )
        .unwrap();
        let dims = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
        let result =
            range_builtin(Value::Tensor(tensor), vec![Value::Tensor(dims)]).expect("range");
        match result {
            Value::Tensor(t) => {
                assert_eq!(t.shape, vec![1, 1, 1]);
                assert_eq!(t.data, vec![11.0]);
            }
            Value::Num(n) => assert_eq!(n, 11.0),
            other => panic!("expected scalar tensor or scalar value, got {other:?}"),
        }
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_invalid_dimension_non_integer() {
        let tensor = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
        let err = range_builtin(Value::Tensor(tensor), vec![Value::Num(1.5)])
            .expect_err("expected dimension error");
        assert!(err.message().contains("integer"));
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_gpu_provider_roundtrip() {
        test_support::with_test_provider(|provider| {
            let tensor = Tensor::new(vec![1.0, 3.0, 4.0, 7.0, 2.0, 5.0], vec![2, 3]).unwrap();
            let view = runmat_accelerate_api::HostTensorView {
                data: &tensor.data,
                shape: &tensor.shape,
            };
            let handle = provider.upload(&view).expect("upload");
            let result = range_builtin(Value::GpuTensor(handle), Vec::new()).expect("range");
            let gathered = test_support::gather(result).expect("gather");
            assert_eq!(gathered.shape, vec![1, 3]);
            assert_eq!(gathered.data, vec![2.0, 3.0, 3.0]);
        });
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_gpu_row_vector_matches_cpu() {
        test_support::with_test_provider(|provider| {
            let tensor = Tensor::new(vec![1.0, 4.0, 7.0], vec![1, 3]).unwrap();
            let view = HostTensorView {
                data: &tensor.data,
                shape: &tensor.shape,
            };
            let handle = provider.upload(&view).expect("upload");
            let result =
                range_builtin(Value::GpuTensor(handle.clone()), Vec::new()).expect("range");
            let gathered = test_support::gather(result).expect("gather");
            assert_eq!(gathered.shape, vec![1, 1]);
            assert_eq!(gathered.data, vec![6.0]);
            let _ = provider.free(&handle);
        });
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_gpu_all_roundtrip() {
        test_support::with_test_provider(|provider| {
            let tensor = Tensor::new(vec![1.0, 6.0, 3.0, 10.0], vec![2, 2]).unwrap();
            let view = HostTensorView {
                data: &tensor.data,
                shape: &tensor.shape,
            };
            let handle = provider.upload(&view).expect("upload");
            let result = range_builtin(Value::GpuTensor(handle.clone()), vec![Value::from("all")])
                .expect("range");
            let gathered = test_support::gather(result).expect("gather");
            assert_eq!(gathered.shape, vec![1, 1]);
            assert_eq!(gathered.data, vec![9.0]);
            let _ = provider.free(&handle);
        });
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    fn range_gpu_omit_nan_falls_back_to_host() {
        test_support::with_test_provider(|provider| {
            let tensor = Tensor::new(vec![2.0, 4.0, f64::NAN, 6.0], vec![2, 2]).unwrap();
            let view = runmat_accelerate_api::HostTensorView {
                data: &tensor.data,
                shape: &tensor.shape,
            };
            let handle = provider.upload(&view).expect("upload");
            let args = vec![Value::from("omitnan")];
            let result = range_builtin(Value::GpuTensor(handle), args).expect("range");
            match result {
                Value::Tensor(t) => {
                    assert_eq!(t.shape, vec![1, 2]);
                    assert_eq!(t.data, vec![2.0, 0.0]);
                }
                other => panic!("expected host tensor fallback, got {other:?}"),
            }
        });
    }

    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
    #[test]
    #[cfg(feature = "wgpu")]
    fn range_wgpu_dim2_matches_cpu() {
        let _ = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
            runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
        );
        let provider = runmat_accelerate_api::provider().expect("wgpu provider");
        let tensor = Tensor::new(vec![1.0, 4.0, 2.0, 6.0, 3.0, 5.0], vec![2, 3]).unwrap();
        let cpu = range_builtin(
            Value::Tensor(tensor.clone()),
            vec![Value::Int(IntValue::I32(2))],
        )
        .expect("cpu range");
        let view = HostTensorView {
            data: &tensor.data,
            shape: &tensor.shape,
        };
        let handle = provider.upload(&view).expect("upload");
        let gpu = range_builtin(
            Value::GpuTensor(handle.clone()),
            vec![Value::Int(IntValue::I32(2))],
        )
        .expect("gpu range");
        let cpu_tensor = test_support::gather(cpu).expect("gather cpu");
        let gpu_tensor = test_support::gather(gpu).expect("gather gpu");
        assert_eq!(gpu_tensor.shape, cpu_tensor.shape);
        let tol = match provider.precision() {
            runmat_accelerate_api::ProviderPrecision::F64 => 1e-12,
            runmat_accelerate_api::ProviderPrecision::F32 => 1e-5,
        };
        for (a, b) in gpu_tensor.data.iter().zip(cpu_tensor.data.iter()) {
            assert!((a - b).abs() < tol, "|{} - {}| >= {}", a, b, tol);
        }
        let _ = provider.free(&handle);
    }
}
