burn-ndarray 0.22.0-pre.3

use alloc::{vec, vec::Vec};

use burn_backend::{
    DType, ExecutionError, Shape, TensorData, TensorMetadata, TensorPrimitive, get_device_settings,
    ops::{FloatTensorOps, QTensorOps},
    quantization::{
        BlockSize, QuantMode, QuantPropagation, QuantScheme, QuantStore, QuantValue,
        QuantizationParametersPrimitive, QuantizedBytes, global_scale_dtype, params_shape,
        scale_to_dtype,
    },
    tensor::{FloatTensor, IntTensor, QuantizedTensor},
};
use burn_std::{FloatDType, IntDType};
use ndarray::ArrayD;

use crate::{
    NdArray, NdArrayDevice, NdArrayQTensor, NdArrayTensor, SharedArray, element::QuantElement,
    execute_with_dtype, execute_with_int_dtype, execute_with_int_out_dtype,
    execute_with_numeric_dtype, slice,
};

use super::quantization::{QuantizationStrategy, SymmetricQuantization};
use super::{NdArrayMathOps, NdArrayOps};

impl QTensorOps<Self> for NdArray {
    fn q_from_data(data: TensorData, _device: &NdArrayDevice) -> QuantizedTensor<Self> {
        match data.dtype {
            DType::QFloat(scheme) => {
                let shape = data.shape.clone();
                let q_bytes = QuantizedBytes {
                    shape: shape.clone(),
                    bytes: data.into_bytes(),
                    scheme,
                };

                match scheme {
                    QuantScheme {
                        mode: QuantMode::Symmetric,
                        value: QuantValue::Q8F | QuantValue::Q8S,
                        ..
                    } => {
                        // We can load QuantStore::U32 w/ QuantizedBytes impl
                        let (values, qparams) = q_bytes.into_vec_i8();
                        let data = TensorData::new(values, shape);
                        // Overwrite storage
                        let scheme = scheme.with_store(QuantStore::Native);

                        let global = qparams.global;
                        let qparams = qparams.block;

                        NdArrayQTensor {
                            qtensor: NdArrayTensor::from_data(data),
                            scheme,
                            qparams,
                            global,
                        }
                    }
                    QuantScheme {
                        value:
                            QuantValue::Q4F
                            | QuantValue::Q4S
                            | QuantValue::Q2F
                            | QuantValue::Q2S
                            | QuantValue::E2M1
                            | QuantValue::E4M3
                            | QuantValue::E5M2,
                        ..
                    }
                    | QuantScheme {
                        mode: QuantMode::Lookup,
                        ..
                    } => unimplemented!("from_data not supported for scheme {scheme:?}"),
                }
            }
            _ => panic!(
                "Invalid dtype (expected DType::QFloat, got {:?})",
                data.dtype
            ),
        }
    }

    fn quantize(
        tensor: FloatTensor<Self>,
        scheme: &QuantScheme,
        qparams: QuantizationParametersPrimitive<Self>,
    ) -> QuantizedTensor<Self> {
        let shape = tensor.shape();
        let data_f = tensor.into_data();
        let scales = qparams.scales.into_data().convert::<f32>();
        // Quantize against the scale that will actually be stored, so a save/load round trip
        // reproduces these values instead of drifting by the scale dtype's rounding error.
        let scales: Vec<f32> = scales
            .iter::<f32>()
            .map(|s| scale_to_dtype(s, scheme.scale_dtype()))
            .collect();
        let global = qparams.global.map(|global| {
            let dtype = global_scale_dtype(scheme)
                .expect("a per-tensor scale should come with a two-level scheme");
            let global = global.into_data().convert::<f32>();
            scale_to_dtype(global.iter::<f32>().next().unwrap(), dtype)
        });

        // Implement with ndarray instead of QuantizationStrategy?
        let (data, qparams) = match (scheme.block_size(), scheme) {
            (
                None,
                QuantScheme {
                    mode: QuantMode::Symmetric,
                    // `Q2S` is supported natively (stored as i8): it feeds the multiply-free
                    // ternary matmul fast path in `q_matmul` (BitNet b1.58).
                    #[cfg(not(feature = "export_tests"))]
                        value: QuantValue::Q8F | QuantValue::Q8S | QuantValue::Q2S,
                    // For tests, "native" sub-byte quant serves as a reference for value equality.
                    // Values are stored as i8 regardless.
                    #[cfg(feature = "export_tests")]
                        value:
                        QuantValue::Q8F
                        | QuantValue::Q8S
                        | QuantValue::Q4F
                        | QuantValue::Q4S
                        | QuantValue::Q2F
                        | QuantValue::Q2S,
                    store: QuantStore::Native,
                    ..
                },
            ) => {
                let scales = scales[0];
                let strategy = QuantizationStrategy::PerTensorSymmetric(
                    SymmetricQuantization::init(scales, scheme.value),
                );
                let values = strategy.quantize(data_f.as_slice().unwrap(), &shape);
                (
                    TensorData::quantized(values, shape.clone(), *scheme, &[scales], None),
                    vec![scales],
                )
            }
            (
                Some(block_size),
                QuantScheme {
                    mode: QuantMode::Symmetric,
                    #[cfg(not(feature = "export_tests"))]
                        value: QuantValue::Q8F | QuantValue::Q8S,
                    #[cfg(feature = "export_tests")]
                        value:
                        QuantValue::Q8F
                        | QuantValue::Q8S
                        | QuantValue::Q4F
                        | QuantValue::Q4S
                        | QuantValue::Q2F
                        | QuantValue::Q2S,
                    store: QuantStore::Native,
                    ..
                },
            ) => {
                let global = if global_scale_dtype(scheme).is_some() {
                    Some(global.expect("a two-level scheme should have a per-tensor scale"))
                } else {
                    None
                };
                quantize_per_block(
                    data_f.as_slice().unwrap(),
                    shape.clone(),
                    scheme,
                    block_size,
                    scales.as_slice(),
                    global,
                )
            }
            (_, scheme) => unimplemented!("Quantization not supported for scheme {scheme:?}"),
        };

        let q_bytes = QuantizedBytes {
            shape: data.shape.clone(),
            bytes: data.into_bytes(),
            scheme: *scheme,
        };
        let (values, _) = q_bytes.into_vec_i8();
        let data = TensorData::new(values, shape);

        NdArrayQTensor {
            qtensor: NdArrayTensor::from_data(data),
            scheme: *scheme,
            qparams,
            global,
        }
    }

    fn dequantize(tensor: QuantizedTensor<Self>, dtype: FloatDType) -> FloatTensor<Self> {
        let strategy = tensor.strategy();
        let scheme = tensor.scheme;
        let shape = tensor.shape();
        let scales = tensor.qparams;
        let global = tensor.global;
        let data = match tensor.qtensor {
            NdArrayTensor::I8(storage) => {
                let data = storage.into_shared().into_iter().collect();
                dequantize(
                    data,
                    shape,
                    scheme,
                    &strategy,
                    &scales,
                    global,
                    dtype.into(),
                )
            }
            _ => unreachable!(),
        };
        NdArrayTensor::from_data(data)
    }

    /// Matrix multiplication with at least one quantized operand.
    ///
    /// Fast path — BitNet b1.58 ternary weights: when `rhs` is a `Q2S` symmetric per-tensor weight
    /// (values in `{-1, 0, +1}`) and `lhs` is an `f32` activation, the product is computed WITHOUT
    /// dequantizing the weight and WITHOUT a single multiply in the inner loop: `+1 => add`,
    /// `-1 => subtract`, `0 => skip`, then the per-tensor scale `γ` is applied once per output
    /// element — the multiply-free compute path BitNet is built on. The result matches the
    /// dequantize-then-`float_matmul` path to within f32 rounding.
    ///
    /// Every other case (Q8, per-block, non-f32 activation, batched weights, ...) falls through to
    /// the regular `dequantize -> float_matmul` path — byte-for-byte the default behaviour.
    fn q_matmul(lhs: TensorPrimitive<Self>, rhs: TensorPrimitive<Self>) -> TensorPrimitive<Self> {
        if let (TensorPrimitive::Float(l), TensorPrimitive::QFloat(r)) = (&lhs, &rhs)
            && let Some(out) = ternary_matmul(l, r)
        {
            return TensorPrimitive::Float(out);
        }

        // Fallback: identical to the default `QTensorOps::q_matmul` — dequantize any quantized
        // operand, run the regular float matmul, and preserve quantization propagation.
        let mut propagation = QuantPropagation::Inhibit;
        let mut scheme = QuantScheme::default();
        let target_dtype: Option<FloatDType> = match (&lhs, &rhs) {
            (TensorPrimitive::Float(t), _) | (_, TensorPrimitive::Float(t)) => {
                Some(t.dtype().into())
            }
            _ => None,
        };
        let lhs = match lhs {
            TensorPrimitive::Float(lhs) => lhs,
            TensorPrimitive::QFloat(lhs) => {
                let settings = get_device_settings::<Self>(&lhs.device());
                propagation = settings.quantization.propagation;
                scheme = lhs.scheme;
                let float_dtype = target_dtype.unwrap_or(settings.float_dtype);
                Self::dequantize(lhs, float_dtype)
            }
        };
        let rhs = match rhs {
            TensorPrimitive::Float(rhs) => rhs,
            TensorPrimitive::QFloat(rhs) => {
                let settings = get_device_settings::<Self>(&rhs.device());
                propagation = settings.quantization.propagation;
                scheme = rhs.scheme;
                let float_dtype = target_dtype.unwrap_or(settings.float_dtype);
                Self::dequantize(rhs, float_dtype)
            }
        };
        let out_f = <Self as FloatTensorOps<Self>>::float_matmul(lhs, rhs);
        match propagation {
            QuantPropagation::Propagate => {
                TensorPrimitive::QFloat(Self::quantize_dynamic(out_f, &scheme))
            }
            QuantPropagation::Inhibit => TensorPrimitive::Float(out_f),
        }
    }

    fn q_to_device(
        tensor: QuantizedTensor<Self>,
        _device: &NdArrayDevice,
    ) -> QuantizedTensor<Self> {
        tensor
    }

    fn q_reshape(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
        NdArrayQTensor {
            qtensor: execute_with_dtype!(tensor.qtensor, E, |array: SharedArray<E>| {
                NdArrayOps::reshape(array, shape)
            }),
            scheme: tensor.scheme,
            qparams: tensor.qparams,
            global: tensor.global,
        }
    }

    async fn q_into_data(tensor: QuantizedTensor<Self>) -> Result<TensorData, ExecutionError> {
        let shape = tensor.qtensor.shape();
        let scales = tensor.qparams;
        Ok(execute_with_numeric_dtype!(
            tensor.qtensor,
            E,
            |array: SharedArray<E>| {
                let values = array.into_iter().collect();
                TensorData::quantized(values, shape, tensor.scheme, &scales, tensor.global)
            }
        ))
    }

    fn q_swap_dims(
        tensor: QuantizedTensor<Self>,
        dim1: usize,
        dim2: usize,
    ) -> QuantizedTensor<Self> {
        let mut axes = (0..tensor.qtensor.shape().num_dims()).collect::<Vec<_>>();
        axes.swap(dim1, dim2);
        Self::q_permute(tensor, &axes)
    }

    fn q_permute(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
        let (scheme, qparams) = match tensor.scheme.block_size() {
            None => (tensor.scheme, tensor.qparams),
            Some(_) => {
                let shape = tensor.qtensor.shape();
                let qparams_shape = params_shape(&shape, &tensor.scheme);
                let scales = ArrayD::from_shape_vec(qparams_shape.as_slice(), tensor.qparams)
                    .unwrap()
                    .into_shared();
                let qparams = NdArrayOps::permute(scales, axes).into_iter().collect();

                let mut scheme = tensor.scheme;
                scheme.permute_block_dims(shape.num_dims(), axes);

                (scheme, qparams)
            }
        };

        NdArrayQTensor {
            qtensor: execute_with_dtype!(tensor.qtensor, E, |array: SharedArray<E>| {
                NdArrayOps::permute(array, axes)
            }),
            scheme,
            qparams,
            global: tensor.global,
        }
    }

    fn q_flip(tensor: QuantizedTensor<Self>, axes: &[usize]) -> QuantizedTensor<Self> {
        NdArrayQTensor {
            qtensor: execute_with_dtype!(tensor.qtensor, E, |array: SharedArray<E>| {
                NdArrayOps::flip(array, axes)
            }),
            scheme: tensor.scheme,
            qparams: tensor.qparams,
            global: tensor.global,
        }
    }

    fn q_gather(
        dim: usize,
        tensor: QuantizedTensor<Self>,
        indices: IntTensor<Self>,
    ) -> QuantizedTensor<Self> {
        let qtensor = execute_with_int_dtype!(indices, IntElem, |idx_array: SharedArray<
            IntElem,
        >|
         -> NdArrayTensor {
            execute_with_numeric_dtype!(tensor.qtensor, E, |array: SharedArray<E>| {
                NdArrayOps::gather(dim, array, idx_array)
            })
        });
        NdArrayQTensor {
            qtensor,
            scheme: tensor.scheme,
            qparams: tensor.qparams,
            global: tensor.global,
        }
    }

    fn q_select(
        tensor: QuantizedTensor<Self>,
        dim: usize,
        indices: IntTensor<Self>,
    ) -> QuantizedTensor<Self> {
        let qtensor = execute_with_int_dtype!(indices, IntElem, |idx_array: SharedArray<
            IntElem,
        >|
         -> NdArrayTensor {
            execute_with_numeric_dtype!(tensor.qtensor, E, |array: SharedArray<E>| {
                NdArrayMathOps::select(array, dim, idx_array)
            })
        });
        NdArrayQTensor {
            qtensor,
            scheme: tensor.scheme,
            qparams: tensor.qparams,
            global: tensor.global,
        }
    }

    fn q_slice(
        tensor: QuantizedTensor<Self>,
        slices: &[burn_backend::Slice],
    ) -> QuantizedTensor<Self> {
        NdArrayQTensor {
            qtensor: slice!(tensor.qtensor, slices),
            scheme: tensor.scheme,
            qparams: tensor.qparams,
            global: tensor.global,
        }
    }

    fn q_argmax(tensor: QuantizedTensor<Self>, dim: usize, out_dtype: IntDType) -> IntTensor<Self> {
        execute_with_int_out_dtype!(out_dtype, I, {
            execute_with_numeric_dtype!(tensor.qtensor, E, |array: SharedArray<E>| {
                NdArrayMathOps::argmax::<I>(array, dim)
            })
        })
    }

    fn q_argmin(tensor: QuantizedTensor<Self>, dim: usize, out_dtype: IntDType) -> IntTensor<Self> {
        execute_with_int_out_dtype!(out_dtype, I, {
            execute_with_numeric_dtype!(tensor.qtensor, E, |array: SharedArray<E>| {
                NdArrayMathOps::argmin::<I>(array, dim)
            })
        })
    }

    fn q_expand(tensor: QuantizedTensor<Self>, shape: Shape) -> QuantizedTensor<Self> {
        NdArrayQTensor {
            qtensor: execute_with_dtype!(tensor.qtensor, E, |array: SharedArray<E>| {
                NdArrayOps::expand(array, shape)
            }),
            scheme: tensor.scheme,
            qparams: tensor.qparams,
            global: tensor.global,
        }
    }
}

/// Native multiply-free ternary matmul (BitNet b1.58): an `f32` activation `lhs` times a `Q2S`
/// symmetric per-tensor ternary weight `rhs` (values in `{-1, 0, +1}`).
///
/// Returns `None` — so the caller falls back to the regular dequantize path — unless every
/// precondition holds: `rhs` is `Q2S` / `Symmetric` / per-tensor, `rhs` is a 2D weight `[K, N]`,
/// and `lhs` is an `f32` tensor whose last dim is `K`. Leading dims of `lhs` are flattened into the
/// row count `M`, so this covers the `Linear`-style `[.., K] x [K, N] -> [.., N]` case.
///
/// The inner loop contains no multiplies: `+1 => add`, `-1 => subtract`, `0 => skip`. The single
/// per-tensor scale `γ` is applied once per output element at the end — `M·N` multiplies instead of
/// the `M·N·K` of a dense matmul, with zeros never touched.
fn ternary_matmul(
    lhs: &FloatTensor<NdArray>,
    rhs: &NdArrayQTensor,
) -> Option<FloatTensor<NdArray>> {
    // Canonical BitNet b1.58 weight quantization: Q2S, symmetric, per-tensor.
    if rhs.scheme.block_size().is_some()
        || !matches!(
            rhs.scheme,
            QuantScheme {
                value: QuantValue::Q2S,
                mode: QuantMode::Symmetric,
                ..
            }
        )
    {
        return None;
    }
    // Only an f32 activation (the reference float dtype) takes the fast path.
    if !matches!(lhs, NdArrayTensor::F32(_)) {
        return None;
    }

    // rhs is a 2D weight [K, N].
    let wdims = rhs.qtensor.shape().to_vec();
    if wdims.len() != 2 {
        return None;
    }
    let (k, n) = (wdims[0], wdims[1]);

    // lhs is [.., M, K] with a matching K; flatten the leading dims into M.
    let ldims = lhs.shape().to_vec();
    if ldims.len() < 2 || *ldims.last().unwrap() != k {
        return None;
    }
    let m: usize = ldims[..ldims.len() - 1].iter().product();

    // Per-tensor scale γ.
    let gamma = *rhs.qparams.first()?;

    // Canonical row-major values for both operands (`into_data` normalizes any strided layout).
    let a_data = lhs.clone().into_data();
    let a = a_data.as_slice::<f32>().ok()?;
    let w_data = rhs.qtensor.clone().into_data();
    let w = w_data.as_slice::<i8>().ok()?;
    if a.len() != m * k || w.len() != k * n {
        return None;
    }

    // Multiply-free accumulation; one scale per output element at the end.
    let mut out = vec![0f32; m * n];
    for i in 0..m {
        let arow = &a[i * k..(i + 1) * k];
        let orow = &mut out[i * n..(i + 1) * n];
        for kk in 0..k {
            let x = arow[kk];
            let wrow = &w[kk * n..(kk + 1) * n];
            for (o, &t) in orow.iter_mut().zip(wrow) {
                match t {
                    1 => *o += x,
                    -1 => *o -= x,
                    _ => {} // 0 -> skip
                }
            }
        }
        for o in orow.iter_mut() {
            *o *= gamma;
        }
    }

    let mut out_dims = ldims[..ldims.len() - 1].to_vec();
    out_dims.push(n);
    Some(NdArrayTensor::from_data(TensorData::new(
        out,
        Shape::from(out_dims),
    )))
}

/// `global: None` is a one-level block scheme (multiplier 1.0).
fn quantize_per_block(
    data_f: &[f32],
    shape: Shape,
    scheme: &QuantScheme,
    block: BlockSize,
    scales: &[f32],
    global: Option<f32>,
) -> (TensorData, Vec<f32>) {
    let multiplier = global.unwrap_or(1.0);
    let (strategy, qparams): (Vec<_>, Vec<_>) = scales
        .iter()
        .map(|&s| (SymmetricQuantization::init(multiplier * s, scheme.value), s))
        .unzip();
    let strategy = QuantizationStrategy::PerBlockSymmetric(strategy, block);
    let values = strategy.quantize(data_f, &shape);
    (
        TensorData::quantized(values, shape, *scheme, scales, global),
        qparams,
    )
}

fn dequantize<Q: QuantElement>(
    data: Vec<Q>,
    shape: Shape,
    scheme: QuantScheme,
    strategy: &QuantizationStrategy,
    qparams: &[f32],
    global: Option<f32>,
    dtype: DType,
) -> TensorData {
    let q_bytes = QuantizedBytes::new(data, shape.clone(), scheme, qparams, global);
    let (values, _qparams) = q_bytes.into_vec_i8();
    let values = strategy.dequantize(&values, &shape);
    TensorData::new(values, shape).convert_dtype(dtype)
}