use std::collections::HashMap;
use std::io::Read;

use burn_flex::{Flex, FlexDevice};
use burn_tensor::activation;
use burn_tensor::backend::Backend;
use burn_tensor::{ElementConversion, Int, Tensor, TensorData};

use crate::backend::{BackendMetadata, InferenceBackend};
use crate::error::{NarsilError, Result};
use crate::schema::{DataType, InferenceRequest, InferenceResponse};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DlrmConfig {
    pub dense_in_features: usize,
    pub sparse_features: usize,
    pub num_embeddings: usize,
    pub embedding_dim: usize,
    pub dense_arch_layer_sizes: Vec<usize>,
    pub over_arch_layer_sizes: Vec<usize>,
}

impl Default for DlrmConfig {
    fn default() -> Self {
        Self {
            dense_in_features: 13,
            sparse_features: 26,
            num_embeddings: 100_000,
            embedding_dim: 64,
            dense_arch_layer_sizes: vec![512, 256, 64],
            over_arch_layer_sizes: vec![512, 512, 256, 1],
        }
    }
}

impl DlrmConfig {
    pub fn validate(&self) -> Result<()> {
        if self.dense_in_features == 0
            || self.sparse_features == 0
            || self.num_embeddings == 0
            || self.embedding_dim == 0
            || self.dense_arch_layer_sizes.is_empty()
            || self.over_arch_layer_sizes.is_empty()
        {
            return Err(NarsilError::InvalidRequest(
                "DLRM dimensions and layer lists must be non-empty".to_string(),
            ));
        }
        if *self.dense_arch_layer_sizes.last().unwrap() != self.embedding_dim {
            return Err(NarsilError::InvalidRequest(format!(
                "DLRM dense arch must end at embedding_dim {}; got {:?}",
                self.embedding_dim, self.dense_arch_layer_sizes
            )));
        }
        Ok(())
    }

    pub fn interaction_input_dim(&self) -> usize {
        let feature_count = self.sparse_features + 1;
        self.embedding_dim + (feature_count * (feature_count - 1) / 2)
    }

    /// Element count of the stacked embedding table, saturating on overflow so the
    /// `Vec::with_capacity` hint can never panic on absurd configs (the real allocation
    /// is driven by the populated vector length, not this hint).
    fn embedding_value_count(&self) -> usize {
        self.sparse_features
            .saturating_mul(self.num_embeddings)
            .saturating_mul(self.embedding_dim)
    }

    /// Guards the `batch * sparse_features` invariant the embedding gather relies on.
    /// Callers that bypass `prepare_request` (e.g. the public `forward` entry points)
    /// use this so a malformed slice yields a clean error instead of an out-of-bounds
    /// panic inside `sparse_embedding_indices`.
    fn check_sparse_len(&self, len: usize, batch: usize) -> Result<()> {
        let expected = batch * self.sparse_features;
        if len != expected {
            return Err(NarsilError::InvalidRequest(format!(
                "DLRM expected {expected} sparse values, got {len}"
            )));
        }
        Ok(())
    }
}

#[derive(Debug, Clone)]
pub struct BurnDlrmModel<B: Backend> {
    config: DlrmConfig,
    device: B::Device,
    embeddings: Tensor<B, 2>,
    interaction_pair_indices: Tensor<B, 1, Int>,
    dense_weights: Vec<Tensor<B, 2>>,
    dense_biases: Vec<Tensor<B, 2>>,
    over_weights: Vec<Tensor<B, 2>>,
    over_biases: Vec<Tensor<B, 2>>,
}

impl<B: Backend> BurnDlrmModel<B> {
    pub fn new(config: DlrmConfig, device: B::Device) -> Result<Self> {
        config.validate()?;
        let mut seed = 17usize;
        let mut embedding_values = Vec::with_capacity(config.embedding_value_count());
        for _ in 0..config.sparse_features {
            let values = deterministic_values(
                config.num_embeddings * config.embedding_dim,
                seed,
                1.0 / (config.embedding_dim as f32).sqrt(),
            );
            embedding_values.extend(values);
            seed += 101;
        }
        let embeddings = tensor_from_f32(
            embedding_values,
            [
                config.sparse_features * config.num_embeddings,
                config.embedding_dim,
            ],
            &device,
        );
        let interaction_pair_indices = tensor_from_i64(
            dlrm_pair_indices(config.sparse_features + 1),
            [config.interaction_input_dim() - config.embedding_dim],
            &device,
        );

        let (dense_weights, dense_biases) = deterministic_layers(
            config.dense_in_features,
            &config.dense_arch_layer_sizes,
            &device,
            &mut seed,
        );
        let (over_weights, over_biases) = deterministic_layers(
            config.interaction_input_dim(),
            &config.over_arch_layer_sizes,
            &device,
            &mut seed,
        );

        Ok(Self {
            config,
            device,
            embeddings,
            interaction_pair_indices,
            dense_weights,
            dense_biases,
            over_weights,
            over_biases,
        })
    }

    pub fn load(path: &std::path::Path, config: DlrmConfig, device: B::Device) -> Result<Self> {
        config.validate()?;
        let mut reader = std::io::BufReader::new(std::fs::File::open(path).map_err(|err| {
            NarsilError::Backend(format!(
                "failed to open DLRM weights from {}: {err}",
                path.display()
            ))
        })?);

        let mut embedding_values = Vec::with_capacity(config.embedding_value_count());
        for _ in 0..config.sparse_features {
            let values =
                read_f32_values(&mut reader, config.num_embeddings * config.embedding_dim)?;
            embedding_values.extend(values);
        }
        let embeddings = tensor_from_f32(
            embedding_values,
            [
                config.sparse_features * config.num_embeddings,
                config.embedding_dim,
            ],
            &device,
        );
        let interaction_pair_indices = tensor_from_i64(
            dlrm_pair_indices(config.sparse_features + 1),
            [config.interaction_input_dim() - config.embedding_dim],
            &device,
        );

        let (dense_weights, dense_biases) = read_layers(
            &mut reader,
            config.dense_in_features,
            &config.dense_arch_layer_sizes,
            &device,
        )?;
        let (over_weights, over_biases) = read_layers(
            &mut reader,
            config.interaction_input_dim(),
            &config.over_arch_layer_sizes,
            &device,
        )?;

        let mut trailing = [0u8; 1];
        if reader.read(&mut trailing).map_err(|err| {
            NarsilError::Backend(format!("failed to check DLRM weight file length: {err}"))
        })? != 0
        {
            return Err(NarsilError::Backend(
                "DLRM weight file has trailing bytes".to_string(),
            ));
        }

        Ok(Self {
            config,
            device,
            embeddings,
            interaction_pair_indices,
            dense_weights,
            dense_biases,
            over_weights,
            over_biases,
        })
    }

    pub fn config(&self) -> &DlrmConfig {
        &self.config
    }

    pub fn forward(
        &self,
        dense_values: Vec<f32>,
        sparse_values: Vec<i64>,
        batch: usize,
    ) -> Result<Vec<f32>> {
        self.config.check_sparse_len(sparse_values.len(), batch)?;
        let dense = tensor_from_f32(
            dense_values,
            [batch, self.config.dense_in_features],
            &self.device,
        );
        let dense_projected = run_mlp(dense, &self.dense_weights, &self.dense_biases, true);

        let over_input = self.interaction_over_input(dense_projected, &sparse_values, batch);
        let logits = run_mlp(over_input, &self.over_weights, &self.over_biases, false);
        let predictions = activation::sigmoid(logits);
        predictions
            .into_data()
            .convert::<f32>()
            .to_vec::<f32>()
            .map_err(|err| {
                NarsilError::Backend(format!("burn DLRM output conversion failed: {err}"))
            })
    }

    fn sparse_embedding_tensor(&self, sparse_values: &[i64], batch: usize) -> Tensor<B, 3> {
        let indices = sparse_embedding_indices(&self.config, sparse_values, batch);
        let indices = tensor_from_i64(indices, [batch * self.config.sparse_features], &self.device);
        self.embeddings.clone().select(0, indices).reshape([
            batch,
            self.config.sparse_features,
            self.config.embedding_dim,
        ])
    }

    fn feature_tensor(
        &self,
        dense_projected: Tensor<B, 2>,
        sparse_values: &[i64],
        batch: usize,
    ) -> Tensor<B, 3> {
        let dense = dense_projected.reshape([batch, 1, self.config.embedding_dim]);
        let sparse = self.sparse_embedding_tensor(sparse_values, batch);
        Tensor::cat(vec![dense, sparse], 1)
    }

    fn interaction_over_input(
        &self,
        dense_projected: Tensor<B, 2>,
        sparse_values: &[i64],
        batch: usize,
    ) -> Tensor<B, 2> {
        let features = self.feature_tensor(dense_projected.clone(), sparse_values, batch);
        let pairwise = self.pairwise_interactions(features, batch);
        Tensor::cat(vec![dense_projected, pairwise], 1)
    }

    /// Batched Gram interaction: `bmm(X, X^T)` then gather the strict upper triangle.
    /// This is generic over `B`, so the CUDA f16 backend reuses the exact same code and
    /// index math as the FP32/Flex path; f16 numeric parity is validated via the
    /// documented TorchRec comparison benchmark rather than a CPU tolerance test.
    fn pairwise_interactions(&self, features: Tensor<B, 3>, batch: usize) -> Tensor<B, 2> {
        let feature_count = self.config.sparse_features + 1;
        let gram = features.clone().matmul(features.swap_dims(1, 2));
        let gram = gram.reshape([batch, feature_count * feature_count]);
        gram.select(1, self.interaction_pair_indices.clone())
    }
}

#[derive(Debug, Clone)]
pub struct BurnDlrmBackend<B: Backend> {
    model: BurnDlrmModel<B>,
    name: &'static str,
    runtime: &'static str,
}

struct PreparedDlrmInput {
    dense_values: Vec<f32>,
    sparse_values: Vec<i64>,
    batch: usize,
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
type BurnCudaF32 = burn_cuda::Cuda<f32, i32>;

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
type CubeCudaF32 = burn_cubecl::CubeBackend<cubecl::cuda::CudaRuntime, f32, i32, u8>;

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
const CUTILE_DLRM_COPY_TILE: usize = 16;

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
#[derive(Clone)]
pub struct CutileDlrmBackend {
    inner: BurnDlrmBackend<BurnCudaF32>,
    stream: std::sync::Arc<cutile::cuda_core::Stream>,
}

impl BurnDlrmBackend<Flex> {
    pub fn flex(config: DlrmConfig) -> Result<Self> {
        Self::new(config, FlexDevice, "burn-flex-dlrm", "burn-flex")
    }

    pub fn flex_load(path: &std::path::Path, config: DlrmConfig) -> Result<Self> {
        Self::load(path, config, FlexDevice, "burn-flex-dlrm", "burn-flex")
    }
}

impl<B: Backend> BurnDlrmBackend<B> {
    pub fn new(
        config: DlrmConfig,
        device: B::Device,
        name: &'static str,
        runtime: &'static str,
    ) -> Result<Self> {
        Ok(Self {
            model: BurnDlrmModel::new(config, device)?,
            name,
            runtime,
        })
    }

    pub fn load(
        path: &std::path::Path,
        config: DlrmConfig,
        device: B::Device,
        name: &'static str,
        runtime: &'static str,
    ) -> Result<Self> {
        Ok(Self {
            model: BurnDlrmModel::load(path, config, device)?,
            name,
            runtime,
        })
    }

    fn prepare_request(&self, request: &InferenceRequest) -> Result<PreparedDlrmInput> {
        let dense = find_tensor(request, &["dense", "float_features"])?;
        let values = find_tensor(request, &["values", "id_list_features.values"])?;
        let lengths = find_tensor(request, &["lengths", "id_list_features.lengths"])?;

        let batch = match dense.shape.as_slice() {
            [features] if *features == self.model.config().dense_in_features => 1,
            [batch, features] if *features == self.model.config().dense_in_features => *batch,
            _ => {
                return Err(NarsilError::InvalidRequest(format!(
                    "DLRM dense input expects shape [{}] or [batch, {}], got {:?}",
                    self.model.config().dense_in_features,
                    self.model.config().dense_in_features,
                    dense.shape
                )));
            }
        };

        let dense_values = dense.f32_values()?;
        let lengths = int_values(lengths)?;
        let sparse_values = int_values(values)?;
        let expected_sparse = batch * self.model.config().sparse_features;

        if lengths.len() != expected_sparse {
            return Err(NarsilError::InvalidRequest(format!(
                "DLRM expected {expected_sparse} sparse lengths, got {}",
                lengths.len()
            )));
        }
        if lengths.iter().any(|value| *value != 1) {
            return Err(NarsilError::InvalidRequest(
                "DLRM benchmark backend currently expects one id per sparse feature".to_string(),
            ));
        }
        self.model
            .config()
            .check_sparse_len(sparse_values.len(), batch)?;

        Ok(PreparedDlrmInput {
            dense_values,
            sparse_values,
            batch,
        })
    }

    fn response_from_output(&self, batch: usize, output: Vec<f32>) -> Result<InferenceResponse> {
        let mut metadata = HashMap::new();
        metadata.insert("backend".to_string(), self.name.to_string());
        metadata.insert("runtime".to_string(), self.runtime.to_string());
        metadata.insert("workload".to_string(), "dlrm".to_string());

        Ok(InferenceResponse {
            outputs: vec![crate::schema::Tensor::from_f32(
                "default",
                vec![batch],
                output,
            )?],
            metadata,
        })
    }

    fn infer_tensor_set(&self, request: InferenceRequest) -> Result<InferenceResponse> {
        let prepared = self.prepare_request(&request)?;
        let output = self.model.forward(
            prepared.dense_values,
            prepared.sparse_values,
            prepared.batch,
        )?;
        self.response_from_output(prepared.batch, output)
    }

    fn infer_batch_optimised(
        &self,
        requests: Vec<InferenceRequest>,
    ) -> Vec<Result<InferenceResponse>> {
        if requests.is_empty() {
            return vec![];
        }

        let mut all_dense_values = Vec::new();
        let mut all_sparse_values = Vec::new();
        let mut batch_sizes = Vec::with_capacity(requests.len());
        let mut errors = Vec::with_capacity(requests.len());

        for request in &requests {
            match self.prepare_request(request) {
                Ok(prepared) => {
                    all_dense_values.extend(prepared.dense_values);
                    all_sparse_values.extend(prepared.sparse_values);
                    batch_sizes.push(prepared.batch);
                    errors.push(None);
                }
                Err(err) => {
                    batch_sizes.push(0);
                    errors.push(Some(err));
                }
            }
        }

        let total_batch: usize = batch_sizes.iter().sum();
        if total_batch == 0 {
            return errors
                .into_iter()
                .map(|err| {
                    Err(err.unwrap_or_else(|| NarsilError::Backend("empty DLRM batch".to_string())))
                })
                .collect();
        }

        let raw_outputs = match self
            .model
            .forward(all_dense_values, all_sparse_values, total_batch)
        {
            Ok(outputs) => outputs,
            Err(err) => {
                return errors
                    .into_iter()
                    .map(|maybe_err| match maybe_err {
                        Some(err) => Err(err),
                        None => Err(NarsilError::Backend(format!(
                            "batched DLRM forward failed: {err}"
                        ))),
                    })
                    .collect();
            }
        };

        let mut offset = 0usize;
        errors
            .into_iter()
            .zip(batch_sizes)
            .map(|(maybe_err, batch)| {
                if let Some(err) = maybe_err {
                    return Err(err);
                }
                let end = offset.checked_add(batch).ok_or_else(|| {
                    NarsilError::Backend("DLRM batch output offset overflow".to_string())
                })?;
                let output = raw_outputs.get(offset..end).ok_or_else(|| {
                    NarsilError::Backend("DLRM batch output slice out of bounds".to_string())
                })?;
                offset = end;
                self.response_from_output(batch, output.to_vec())
            })
            .collect()
    }
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
impl BurnDlrmModel<BurnCudaF32> {
    pub fn forward_cutile_interaction(
        &self,
        dense_values: Vec<f32>,
        sparse_values: Vec<i64>,
        batch: usize,
        stream: &std::sync::Arc<cutile::cuda_core::Stream>,
    ) -> Result<Vec<f32>> {
        if !self
            .config
            .embedding_dim
            .is_multiple_of(CUTILE_DLRM_COPY_TILE)
        {
            return Err(NarsilError::Backend(format!(
                "cuTile DLRM interaction requires embedding_dim divisible by {CUTILE_DLRM_COPY_TILE}; got {}",
                self.config.embedding_dim
            )));
        }
        self.config.check_sparse_len(sparse_values.len(), batch)?;

        let dense = tensor_from_f32(
            dense_values,
            [batch, self.config.dense_in_features],
            &self.device,
        );
        let dense_projected = run_mlp(dense, &self.dense_weights, &self.dense_biases, true);

        let feature_tensor = self.feature_tensor(dense_projected, &sparse_values, batch);
        let feature_tensor = resolve_fusion_tensor_to_cube(feature_tensor);
        let over_input = run_cutile_dlrm_interaction(
            feature_tensor,
            batch,
            self.config.sparse_features + 1,
            self.config.embedding_dim,
            self.config.interaction_input_dim(),
            self.device.index,
            stream,
        )?;

        let over_input = cube_tensor_to_burn::<2>(over_input);
        let over_weights = self
            .over_weights
            .iter()
            .cloned()
            .map(resolve_fusion_tensor_to_cube)
            .map(cube_tensor_to_burn::<2>)
            .collect::<Vec<_>>();
        let over_biases = self
            .over_biases
            .iter()
            .cloned()
            .map(resolve_fusion_tensor_to_cube)
            .map(cube_tensor_to_burn::<2>)
            .collect::<Vec<_>>();

        let logits = run_mlp(over_input, &over_weights, &over_biases, false);
        let predictions = activation::sigmoid(logits);
        predictions
            .into_data()
            .convert::<f32>()
            .to_vec::<f32>()
            .map_err(|err| {
                NarsilError::Backend(format!("burn DLRM cuTile output conversion failed: {err}"))
            })
    }
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
impl CutileDlrmBackend {
    pub fn new(
        config: DlrmConfig,
        device: burn_cuda::CudaDevice,
        name: &'static str,
        runtime: &'static str,
    ) -> Result<Self> {
        let stream = cutile_stream(device.index)?;
        Ok(Self {
            inner: BurnDlrmBackend::new(config, device, name, runtime)?,
            stream,
        })
    }

    pub fn load(
        path: &std::path::Path,
        config: DlrmConfig,
        device: burn_cuda::CudaDevice,
        name: &'static str,
        runtime: &'static str,
    ) -> Result<Self> {
        let stream = cutile_stream(device.index)?;
        Ok(Self {
            inner: BurnDlrmBackend::load(path, config, device, name, runtime)?,
            stream,
        })
    }

    fn infer_tensor_set(&self, request: InferenceRequest) -> Result<InferenceResponse> {
        let prepared = self.inner.prepare_request(&request)?;
        let output = self.inner.model.forward_cutile_interaction(
            prepared.dense_values,
            prepared.sparse_values,
            prepared.batch,
            &self.stream,
        )?;
        self.inner.response_from_output(prepared.batch, output)
    }

    fn infer_batch_optimised(
        &self,
        requests: Vec<InferenceRequest>,
    ) -> Vec<Result<InferenceResponse>> {
        if requests.is_empty() {
            return vec![];
        }

        let mut all_dense_values = Vec::new();
        let mut all_sparse_values = Vec::new();
        let mut batch_sizes = Vec::with_capacity(requests.len());
        let mut errors = Vec::with_capacity(requests.len());

        for request in &requests {
            match self.inner.prepare_request(request) {
                Ok(prepared) => {
                    all_dense_values.extend(prepared.dense_values);
                    all_sparse_values.extend(prepared.sparse_values);
                    batch_sizes.push(prepared.batch);
                    errors.push(None);
                }
                Err(err) => {
                    batch_sizes.push(0);
                    errors.push(Some(err));
                }
            }
        }

        let total_batch: usize = batch_sizes.iter().sum();
        if total_batch == 0 {
            return errors
                .into_iter()
                .map(|err| {
                    Err(err.unwrap_or_else(|| NarsilError::Backend("empty DLRM batch".to_string())))
                })
                .collect();
        }

        let raw_outputs = match self.inner.model.forward_cutile_interaction(
            all_dense_values,
            all_sparse_values,
            total_batch,
            &self.stream,
        ) {
            Ok(outputs) => outputs,
            Err(err) => {
                return errors
                    .into_iter()
                    .map(|maybe_err| match maybe_err {
                        Some(err) => Err(err),
                        None => Err(NarsilError::Backend(format!(
                            "batched cuTile DLRM forward failed: {err}"
                        ))),
                    })
                    .collect();
            }
        };

        let mut offset = 0usize;
        errors
            .into_iter()
            .zip(batch_sizes)
            .map(|(maybe_err, batch)| {
                if let Some(err) = maybe_err {
                    return Err(err);
                }
                let end = offset.checked_add(batch).ok_or_else(|| {
                    NarsilError::Backend("DLRM batch output offset overflow".to_string())
                })?;
                let output = raw_outputs.get(offset..end).ok_or_else(|| {
                    NarsilError::Backend("DLRM batch output slice out of bounds".to_string())
                })?;
                offset = end;
                self.inner.response_from_output(batch, output.to_vec())
            })
            .collect()
    }
}

impl<B: Backend> InferenceBackend for BurnDlrmBackend<B> {
    fn metadata(&self) -> BackendMetadata {
        BackendMetadata {
            name: self.name.to_string(),
            kind: self.runtime.to_string(),
        }
    }

    fn infer(&self, request: InferenceRequest) -> Result<InferenceResponse> {
        self.infer_tensor_set(request)
    }

    fn infer_batch(&self, requests: Vec<InferenceRequest>) -> Vec<Result<InferenceResponse>> {
        self.infer_batch_optimised(requests)
    }
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
impl InferenceBackend for CutileDlrmBackend {
    fn metadata(&self) -> BackendMetadata {
        self.inner.metadata()
    }

    fn infer(&self, request: InferenceRequest) -> Result<InferenceResponse> {
        self.infer_tensor_set(request)
    }

    fn infer_batch(&self, requests: Vec<InferenceRequest>) -> Vec<Result<InferenceResponse>> {
        self.infer_batch_optimised(requests)
    }
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
fn cutile_stream(device_index: usize) -> Result<std::sync::Arc<cutile::cuda_core::Stream>> {
    let device = cutile::cuda_core::Device::new(device_index).map_err(|err| {
        NarsilError::Backend(format!(
            "failed to initialize cuTile CUDA device {device_index}: {err}"
        ))
    })?;
    device.new_stream().map_err(|err| {
        NarsilError::Backend(format!(
            "failed to create cuTile CUDA stream for device {device_index}: {err}"
        ))
    })
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
fn resolve_fusion_tensor_to_cube<const D: usize>(
    tensor: Tensor<BurnCudaF32, D>,
) -> burn_cubecl::tensor::CubeTensor<cubecl::cuda::CudaRuntime> {
    let primitive = tensor.into_primitive().tensor();
    let client = primitive.client.clone();
    client.resolve_tensor_float::<CubeCudaF32>(primitive)
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
fn cube_tensor_to_burn<const D: usize>(
    tensor: burn_cubecl::tensor::CubeTensor<cubecl::cuda::CudaRuntime>,
) -> Tensor<CubeCudaF32, D> {
    Tensor::from_primitive(burn_tensor::backend::TensorPrimitive::Float(tensor))
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
fn cube_device_pointer(
    tensor: &burn_cubecl::tensor::CubeTensor<cubecl::cuda::CudaRuntime>,
) -> Result<cutile::tile_kernel::DevicePointer<f32>> {
    cubecl_common::future::block_on(tensor.client.sync()).map_err(|err| {
        NarsilError::Backend(format!(
            "failed to synchronize CubeCL before cuTile launch: {err}"
        ))
    })?;
    let resource = tensor
        .client
        .get_resource(tensor.handle.clone())
        .map_err(|err| {
            NarsilError::Backend(format!(
                "failed to get CubeCL CUDA resource for cuTile: {err}"
            ))
        })?;
    let pointer = resource.resource().ptr;
    Ok(unsafe { cutile::tile_kernel::DevicePointer::from_cu_deviceptr(pointer) })
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
fn checked_i32(value: usize, name: &str) -> Result<i32> {
    i32::try_from(value)
        .map_err(|_| NarsilError::Backend(format!("cuTile DLRM {name} exceeds i32 range: {value}")))
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
fn run_cutile_dlrm_interaction(
    features: burn_cubecl::tensor::CubeTensor<cubecl::cuda::CudaRuntime>,
    batch: usize,
    feature_count: usize,
    embedding_dim: usize,
    output_dim: usize,
    device_index: usize,
    stream: &std::sync::Arc<cutile::cuda_core::Stream>,
) -> Result<burn_cubecl::tensor::CubeTensor<cubecl::cuda::CudaRuntime>> {
    use cutile::tile_kernel::{DeviceOp, TileKernel};

    let pair_count = feature_count
        .checked_mul(feature_count.saturating_sub(1))
        .and_then(|count| count.checked_div(2))
        .ok_or_else(|| NarsilError::Backend("cuTile DLRM pair count overflow".to_string()))?;
    let dense_tiles = embedding_dim / CUTILE_DLRM_COPY_TILE;
    let work_items = dense_tiles
        .checked_add(pair_count)
        .ok_or_else(|| NarsilError::Backend("cuTile DLRM work-item count overflow".to_string()))?;

    let output = Tensor::<CubeCudaF32, 2>::empty([batch, output_dim], &features.device)
        .into_primitive()
        .tensor();
    let features_ptr = cube_device_pointer(&features)?;
    let output_ptr = cube_device_pointer(&output)?;

    let batch_i32 = checked_i32(batch, "batch")?;
    let feature_count_i32 = checked_i32(feature_count, "feature_count")?;
    let pair_count_i32 = checked_i32(pair_count, "pair_count")?;
    let output_dim_i32 = checked_i32(output_dim, "output_dim")?;
    let grid = (
        u32::try_from(batch)
            .map_err(|_| NarsilError::Backend(format!("cuTile DLRM batch too large: {batch}")))?,
        u32::try_from(work_items).map_err(|_| {
            NarsilError::Backend(format!(
                "cuTile DLRM work item count too large: {work_items}"
            ))
        })?,
        1,
    );
    let generics = vec![embedding_dim.to_string(), CUTILE_DLRM_COPY_TILE.to_string()];

    let launch = unsafe {
        cutile_dlrm_kernel::dlrm_interaction(
            output_ptr,
            features_ptr,
            batch_i32,
            feature_count_i32,
            pair_count_i32,
            output_dim_i32,
        )
    }
    .grid(grid)
    .generics(generics);

    let (_output_ptr, _features_ptr, _batch, _features, _pairs, _output_dim) =
        launch.sync_on(stream).map_err(|err| {
            NarsilError::Backend(format!(
                "cuTile DLRM interaction launch failed on CUDA device {device_index}: {err}"
            ))
        })?;

    Ok(output)
}

#[cfg(all(feature = "cuda", feature = "cuda-cutile"))]
#[cutile::module]
mod cutile_dlrm_kernel {
    use cutile::core::*;

    unsafe fn tensor2_from_ptr<T: ElementType>(
        ptr: *mut T,
        dim0: i32,
        dim1: i32,
    ) -> Tensor<T, { [-1, -1] }> {
        let shape: Shape<{ [-1, -1] }> = Shape::<{ [-1, -1] }> {
            dims: &[dim0, dim1],
        };
        let strides: Array<{ [-1, -1] }> = Array::<{ [-1, -1] }> {
            dims: &[dim1, 1i32],
        };
        let ptr_tile: PointerTile<*mut T, { [] }> = pointer_to_tile(ptr);
        unsafe { make_tensor_view(ptr_tile, shape, strides, new_token_unordered()) }
    }

    unsafe fn tensor3_from_ptr<T: ElementType>(
        ptr: *mut T,
        dim0: i32,
        dim1: i32,
        dim2: i32,
    ) -> Tensor<T, { [-1, -1, -1] }> {
        let shape: Shape<{ [-1, -1, -1] }> = Shape::<{ [-1, -1, -1] }> {
            dims: &[dim0, dim1, dim2],
        };
        let strides: Array<{ [-1, -1, -1] }> = Array::<{ [-1, -1, -1] }> {
            dims: &[dim1 * dim2, dim2, 1i32],
        };
        let ptr_tile: PointerTile<*mut T, { [] }> = pointer_to_tile(ptr);
        unsafe { make_tensor_view(ptr_tile, shape, strides, new_token_unordered()) }
    }

    #[cutile::entry()]
    unsafe fn dlrm_interaction<const EMBEDDING_DIM: i32, const COPY_TILE: i32>(
        output_ptr: *mut f32,
        features_ptr: *mut f32,
        batch: i32,
        feature_count: i32,
        pair_count: i32,
        output_dim: i32,
    ) {
        let pid: (i32, i32, i32) = get_tile_block_id();
        let batch_idx = pid.0;
        let work_idx = pid.1;
        let dense_tiles = EMBEDDING_DIM / COPY_TILE;
        let mut output: Tensor<f32, { [-1, -1] }> =
            unsafe { tensor2_from_ptr(output_ptr, batch, output_dim) };
        let features: Tensor<f32, { [-1, -1, -1] }> =
            unsafe { tensor3_from_ptr(features_ptr, batch, feature_count, EMBEDDING_DIM) };

        if work_idx < dense_tiles {
            let feature_part = features.partition(const_shape![1, 1, COPY_TILE]);
            let dense_tile: Tile<f32, { [1, COPY_TILE] }> = feature_part
                .load([batch_idx, 0i32, work_idx])
                .reshape(const_shape![1, COPY_TILE]);
            let mut output_part: PartitionMut<f32, { [1, COPY_TILE] }> =
                unsafe { output.partition_mut(const_shape![1, COPY_TILE]) };
            unsafe { output_part.store(dense_tile, [batch_idx, work_idx]) };
        } else {
            let pair_idx = work_idx - dense_tiles;
            if pair_idx < pair_count {
                let mut left = 1i32;
                let mut right = pair_idx;
                while right >= left {
                    right -= left;
                    left += 1i32;
                }

                let feature_part = features.partition(const_shape![1, 1, EMBEDDING_DIM]);
                let left_tile: Tile<f32, { [EMBEDDING_DIM] }> = feature_part
                    .load([batch_idx, left, 0i32])
                    .reshape(const_shape![EMBEDDING_DIM]);
                let right_tile: Tile<f32, { [EMBEDDING_DIM] }> = feature_part
                    .load([batch_idx, right, 0i32])
                    .reshape(const_shape![EMBEDDING_DIM]);
                let product: Tile<f32, { [1, EMBEDDING_DIM] }> =
                    (left_tile * right_tile).reshape(const_shape![1, EMBEDDING_DIM]);
                let dot: Tile<f32, { [1] }> = reduce_sum(product, 1i32);
                let dot: Tile<f32, { [] }> = dot.reshape(const_shape![]);
                let out_col = EMBEDDING_DIM + pair_idx;
                let mut output_part: PartitionMut<f32, { [1, 1] }> =
                    unsafe { output.partition_mut(const_shape![1, 1]) };
                unsafe { output_part.store(dot.reshape(const_shape![1, 1]), [batch_idx, out_col]) };
            }
        }
    }
}

pub(crate) fn find_tensor<'a>(
    request: &'a InferenceRequest,
    names: &[&str],
) -> Result<&'a crate::schema::Tensor> {
    request
        .inputs
        .iter()
        .find(|tensor| names.iter().any(|name| tensor.name == *name))
        .ok_or_else(|| {
            NarsilError::InvalidRequest(format!(
                "DLRM request missing tensor named one of {:?}",
                names
            ))
        })
}

pub(crate) fn int_values(tensor: &crate::schema::Tensor) -> Result<Vec<i64>> {
    match tensor.dtype {
        DataType::I32 => Ok(tensor
            .data
            .chunks_exact(4)
            .map(|chunk| i32::from_le_bytes(chunk.try_into().unwrap()) as i64)
            .collect()),
        DataType::I64 => Ok(tensor
            .data
            .chunks_exact(8)
            .map(|chunk| i64::from_le_bytes(chunk.try_into().unwrap()))
            .collect()),
        dtype => Err(NarsilError::Schema(format!(
            "tensor '{}' must be I32 or I64, got {:?}",
            tensor.name, dtype
        ))),
    }
}

fn run_mlp<B: Backend>(
    mut input: Tensor<B, 2>,
    weights: &[Tensor<B, 2>],
    biases: &[Tensor<B, 2>],
    relu_last: bool,
) -> Tensor<B, 2> {
    for (index, (weight, bias)) in weights.iter().zip(biases.iter()).enumerate() {
        input = input.matmul(weight.clone()) + bias.clone();
        if relu_last || index + 1 != weights.len() {
            input = activation::relu(input);
        }
    }
    input
}

/// Per-layer (weights, biases) for a stack of dense layers.
type LayerParams<B> = (Vec<Tensor<B, 2>>, Vec<Tensor<B, 2>>);

fn deterministic_layers<B: Backend>(
    mut input_dim: usize,
    layer_sizes: &[usize],
    device: &B::Device,
    seed: &mut usize,
) -> LayerParams<B> {
    let mut weights = Vec::with_capacity(layer_sizes.len());
    let mut biases = Vec::with_capacity(layer_sizes.len());
    for &output_dim in layer_sizes {
        let scale = 1.0 / (input_dim as f32).sqrt();
        weights.push(tensor_from_f32(
            deterministic_values(input_dim * output_dim, *seed, scale),
            [input_dim, output_dim],
            device,
        ));
        *seed += 101;
        biases.push(tensor_from_f32(
            deterministic_values(output_dim, *seed, 0.01),
            [1, output_dim],
            device,
        ));
        *seed += 101;
        input_dim = output_dim;
    }
    (weights, biases)
}

fn read_layers<B: Backend>(
    reader: &mut impl Read,
    mut input_dim: usize,
    layer_sizes: &[usize],
    device: &B::Device,
) -> Result<LayerParams<B>> {
    let mut weights = Vec::with_capacity(layer_sizes.len());
    let mut biases = Vec::with_capacity(layer_sizes.len());
    for &output_dim in layer_sizes {
        weights.push(tensor_from_f32(
            read_f32_values(reader, input_dim * output_dim)?,
            [input_dim, output_dim],
            device,
        ));
        biases.push(tensor_from_f32(
            read_f32_values(reader, output_dim)?,
            [1, output_dim],
            device,
        ));
        input_dim = output_dim;
    }
    Ok((weights, biases))
}

fn tensor_from_f32<B: Backend, const D: usize>(
    values: Vec<f32>,
    shape: [usize; D],
    device: &B::Device,
) -> Tensor<B, D> {
    let values = values
        .into_iter()
        .map(B::FloatElem::from_elem)
        .collect::<Vec<_>>();
    Tensor::<B, D>::from_data(TensorData::new(values, shape), device)
}

fn tensor_from_i64<B: Backend, const D: usize>(
    values: Vec<i64>,
    shape: [usize; D],
    device: &B::Device,
) -> Tensor<B, D, Int> {
    Tensor::<B, D, Int>::from_data(TensorData::new(values, shape), device)
}

fn sparse_embedding_indices(config: &DlrmConfig, sparse_values: &[i64], batch: usize) -> Vec<i64> {
    let rows_per_table = config.num_embeddings as i64;
    let mut indices = Vec::with_capacity(batch * config.sparse_features);
    for batch_idx in 0..batch {
        let sample_offset = batch_idx * config.sparse_features;
        for table_idx in 0..config.sparse_features {
            let table_offset = (table_idx * config.num_embeddings) as i64;
            let row = sparse_values[sample_offset + table_idx].rem_euclid(rows_per_table);
            indices.push(table_offset + row);
        }
    }
    indices
}

fn dlrm_pair_indices(feature_count: usize) -> Vec<i64> {
    // `feature_count.saturating_sub(1)` avoids an underflow panic for `feature_count == 0`;
    // the loop below produces zero pairs in that case anyway. Capacity is only a hint.
    let capacity = feature_count.saturating_mul(feature_count.saturating_sub(1)) / 2;
    let mut indices = Vec::with_capacity(capacity);
    for left in 1..feature_count {
        for right in 0..left {
            indices.push((right * feature_count + left) as i64);
        }
    }
    indices
}

fn read_f32_values(reader: &mut impl Read, count: usize) -> Result<Vec<f32>> {
    let byte_count = count
        .checked_mul(4)
        .ok_or_else(|| NarsilError::Backend("DLRM weight byte count overflow".to_string()))?;
    let mut bytes = vec![0u8; byte_count];
    reader.read_exact(&mut bytes).map_err(|err| {
        NarsilError::Backend(format!(
            "failed to read {byte_count} DLRM weight bytes: {err}"
        ))
    })?;
    Ok(bytes
        .chunks_exact(4)
        .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
        .collect())
}

fn deterministic_values(count: usize, seed: usize, scale: f32) -> Vec<f32> {
    (0..count)
        .map(|index| {
            let value = ((index.wrapping_mul(37) + seed.wrapping_mul(101) + 17) % 1009) as f32;
            ((value / 1009.0) - 0.5) * scale
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn small_config() -> DlrmConfig {
        DlrmConfig {
            dense_in_features: 4,
            sparse_features: 3,
            num_embeddings: 11,
            embedding_dim: 4,
            dense_arch_layer_sizes: vec![5, 4],
            over_arch_layer_sizes: vec![8, 1],
        }
    }

    fn forward_legacy<B: Backend>(
        model: &BurnDlrmModel<B>,
        dense_values: Vec<f32>,
        sparse_values: Vec<i64>,
        batch: usize,
    ) -> Result<Vec<f32>> {
        let dense = tensor_from_f32(
            dense_values,
            [batch, model.config.dense_in_features],
            &model.device,
        );
        let dense_projected = run_mlp(dense, &model.dense_weights, &model.dense_biases, true);

        let mut features = Vec::with_capacity(model.config.sparse_features + 1);
        features.push(dense_projected.clone());

        for table_idx in 0..model.config.sparse_features {
            let table_offset = table_idx * model.config.num_embeddings;
            let indices = (0..batch)
                .map(|batch_idx| {
                    let sparse_index = batch_idx * model.config.sparse_features + table_idx;
                    let row =
                        sparse_values[sparse_index].rem_euclid(model.config.num_embeddings as i64);
                    table_offset as i64 + row
                })
                .collect::<Vec<_>>();
            let indices = tensor_from_i64(indices, [batch], &model.device);
            features.push(model.embeddings.clone().select(0, indices));
        }

        let mut pairwise = Vec::with_capacity(features.len() * (features.len() - 1) / 2);
        for left in 0..features.len() {
            for right in 0..left {
                pairwise.push((features[left].clone() * features[right].clone()).sum_dim(1));
            }
        }

        let pairwise = Tensor::cat(pairwise, 1);
        let over_input = Tensor::cat(vec![dense_projected, pairwise], 1);
        let logits = run_mlp(over_input, &model.over_weights, &model.over_biases, false);
        activation::sigmoid(logits)
            .into_data()
            .convert::<f32>()
            .to_vec::<f32>()
            .map_err(|err| NarsilError::Backend(format!("legacy output conversion failed: {err}")))
    }

    fn assert_close(left: &[f32], right: &[f32], tolerance: f32) {
        assert_eq!(left.len(), right.len());
        for (index, (left, right)) in left.iter().zip(right.iter()).enumerate() {
            let delta = (left - right).abs();
            assert!(
                delta <= tolerance,
                "mismatch at index {index}: {left} vs {right}, delta {delta}"
            );
        }
    }

    #[test]
    fn sparse_embedding_indices_are_sample_major_with_table_offsets() {
        let config = small_config();
        let indices = sparse_embedding_indices(&config, &[0, 10, 11, -1, 25, -13], 2);
        assert_eq!(indices, vec![0, 21, 22, 10, 14, 31]);
    }

    #[test]
    fn pair_indices_use_strict_upper_triangle_in_legacy_order() {
        assert_eq!(dlrm_pair_indices(4), vec![1, 2, 6, 3, 7, 11]);
    }

    /// Locks the gram-flatten index math for an additional feature count. The legacy
    /// pairwise loop visits `(left, right)` with `right < left`; the batched path reads
    /// those same dot products from the row-major flattened `F * F` gram matrix at
    /// `right * F + left`. This reconstructs the expected indices independently and also
    /// checks the count equals `F * (F - 1) / 2`, covering the f16 path too since it
    /// shares this generic index helper (the f16 numeric path is validated via the
    /// documented benchmark, not a CPU tolerance test, as CUDA is unavailable here).
    #[test]
    fn pair_indices_match_flattened_gram_upper_triangle_for_f5() {
        let feature_count = 5;
        let mut expected = Vec::new();
        for left in 0..feature_count {
            for right in 0..left {
                expected.push((right * feature_count + left) as i64);
            }
        }
        let indices = dlrm_pair_indices(feature_count);
        assert_eq!(indices, expected);
        assert_eq!(indices.len(), feature_count * (feature_count - 1) / 2);
    }

    #[test]
    fn pair_indices_empty_for_degenerate_feature_counts() {
        assert!(dlrm_pair_indices(0).is_empty());
        assert!(dlrm_pair_indices(1).is_empty());
    }

    #[test]
    fn forward_rejects_mismatched_sparse_length() {
        let model = BurnDlrmModel::<Flex>::new(small_config(), FlexDevice).unwrap();
        // batch=2 expects 2 * sparse_features (3) = 6 sparse values; supply 5.
        let dense = vec![0.1, 0.2, 0.3, 0.4, 0.4, 0.3, 0.2, 0.1];
        let err = model.forward(dense, vec![1, 2, 3, 4, 5], 2).unwrap_err();
        assert!(matches!(err, NarsilError::InvalidRequest(_)));
    }

    #[test]
    fn batched_interaction_forward_matches_legacy_pairwise_path() {
        let model = BurnDlrmModel::<Flex>::new(small_config(), FlexDevice).unwrap();
        let dense = vec![
            0.10, 0.20, 0.30, 0.40, 0.40, 0.30, 0.20, 0.10, -0.20, 0.15, 0.35, -0.05,
        ];
        let sparse = vec![1, 2, 3, 4, 5, 6, -1, 11, 25];

        let expected = forward_legacy(&model, dense.clone(), sparse.clone(), 3).unwrap();
        let actual = model.forward(dense, sparse, 3).unwrap();

        assert_close(&actual, &expected, 1e-5);
    }
}
