// Copyright (c) 2024 RustyTorch++ Team
// Licensed under the Apache License, Version 2.0

//! Sparse matrix implementations for finite element assembly.

use crate::error::{AssemblyError, FeaResult};
use nalgebra::DVector;
// use rtx_tensor::Tensor; // Optional dependency
use std::collections::HashMap;

/// Compressed Sparse Row (CSR) matrix format.
#[derive(Debug, Clone)]
pub struct SparseMatrix {
    /// Number of rows
    nrows: usize,
    /// Number of columns
    ncols: usize,
    /// Row pointers
    row_ptr: Vec<usize>,
    /// Column indices
    col_idx: Vec<usize>,
    /// Non-zero values
    values: Vec<f64>,
    /// Assembly map for efficient insertion during assembly
    assembly_map: HashMap<(usize, usize), usize>,
    /// Whether the matrix has been finalized
    finalized: bool,
}

impl SparseMatrix {
    /// Create a new sparse matrix.
    pub fn new(nrows: usize, ncols: usize) -> Self {
        Self {
            nrows,
            ncols,
            row_ptr: vec![0; nrows + 1],
            col_idx: Vec::new(),
            values: Vec::new(),
            assembly_map: HashMap::new(),
            finalized: false,
        }
    }

    /// Create from triplet format (row, col, value).
    pub fn from_triplets(
        nrows: usize,
        ncols: usize,
        triplets: &[(usize, usize, f64)],
    ) -> FeaResult<Self> {
        let mut matrix = Self::new(nrows, ncols);

        for &(row, col, value) in triplets {
            matrix.add_entry(row, col, value)?;
        }

        matrix.finalize()?;
        Ok(matrix)
    }

    /// Get number of rows.
    pub fn nrows(&self) -> usize {
        self.nrows
    }

    /// Get number of columns.
    pub fn ncols(&self) -> usize {
        self.ncols
    }

    /// Get number of non-zeros.
    pub fn nnz(&self) -> usize {
        self.values.len()
    }

    /// Get matrix density.
    pub fn density(&self) -> f64 {
        if self.nrows == 0 || self.ncols == 0 {
            0.0
        } else {
            self.nnz() as f64 / (self.nrows * self.ncols) as f64
        }
    }

    /// Get matrix bandwidth.
    pub fn bandwidth(&self) -> usize {
        if !self.finalized {
            // For unfinalized matrix, compute from assembly_map
            let mut max_bandwidth = 0;
            for (row, col) in self.assembly_map.keys() {
                let distance = row.abs_diff(*col);
                max_bandwidth = max_bandwidth.max(distance);
            }
            max_bandwidth
        } else {
            // For finalized matrix, compute from CSR structure
            let mut max_bandwidth = 0;
            for row in 0..self.nrows {
                let start = self.row_ptr[row];
                let end = self.row_ptr[row + 1];
                for idx in start..end {
                    let col = self.col_idx[idx];
                    let distance = row.abs_diff(col);
                    max_bandwidth = max_bandwidth.max(distance);
                }
            }
            max_bandwidth
        }
    }

    /// Add entry to the matrix (accumulates if entry exists).
    pub fn add_entry(&mut self, row: usize, col: usize, value: f64) -> FeaResult<()> {
        if self.finalized {
            return Err(AssemblyError::SparseAssemblyFailed {
                reason: "Cannot add entries to finalized matrix".to_string(),
            }
            .into());
        }

        if row >= self.nrows || col >= self.ncols {
            return Err(AssemblyError::MatrixDimensionMismatch {
                expected_rows: self.nrows,
                expected_cols: self.ncols,
                actual_rows: row + 1,
                actual_cols: col + 1,
            }
            .into());
        }

        if value.abs() < 1e-15 {
            return Ok(()); // Skip near-zero entries
        }

        let key = (row, col);
        if let Some(&idx) = self.assembly_map.get(&key) {
            // Entry exists, accumulate
            self.values[idx] += value;
        } else {
            // New entry
            let idx = self.values.len();
            self.values.push(value);
            self.col_idx.push(col);
            self.assembly_map.insert(key, idx);
        }

        Ok(())
    }

    /// Set entry in the matrix (overwrites if entry exists).
    pub fn set_entry(&mut self, row: usize, col: usize, value: f64) -> FeaResult<()> {
        if self.finalized {
            return Err(AssemblyError::SparseAssemblyFailed {
                reason: "Cannot set entries in finalized matrix".to_string(),
            }
            .into());
        }

        let key = (row, col);
        if let Some(&idx) = self.assembly_map.get(&key) {
            // Entry exists, overwrite
            self.values[idx] = value;
        } else if value.abs() >= 1e-15 {
            // New non-zero entry
            let idx = self.values.len();
            self.values.push(value);
            self.col_idx.push(col);
            self.assembly_map.insert(key, idx);
        }

        Ok(())
    }

    /// Apply Dirichlet boundary condition.
    pub fn set_dirichlet(&mut self, dof: usize, _value: f64) -> FeaResult<()> {
        if dof >= self.nrows {
            return Err(AssemblyError::GlobalDofOutOfBounds {
                index: dof,
                max_index: self.nrows.saturating_sub(1),
            }
            .into());
        }

        // Zero out row and column
        for col in 0..self.ncols {
            if col != dof {
                self.set_entry(dof, col, 0.0)?;
            }
        }

        for row in 0..self.nrows {
            if row != dof {
                self.set_entry(row, dof, 0.0)?;
            }
        }

        // Set diagonal entry to 1
        self.set_entry(dof, dof, 1.0)?;

        Ok(())
    }

    /// Finalize the matrix (convert to CSR format).
    pub fn finalize(&mut self) -> FeaResult<()> {
        if self.finalized {
            return Ok(());
        }

        // Group entries by row
        let mut row_entries: Vec<Vec<(usize, f64)>> = vec![Vec::new(); self.nrows];

        for (&(row, col), &idx) in &self.assembly_map {
            let value = self.values[idx];
            if value.abs() >= 1e-15 {
                row_entries[row].push((col, value));
            }
        }

        // Sort columns within each row
        for row_data in &mut row_entries {
            row_data.sort_by_key(|&(col, _)| col);
        }

        // Build CSR format
        self.values.clear();
        self.col_idx.clear();
        self.row_ptr = vec![0; self.nrows + 1];

        let mut nnz = 0;
        for (row, row_data) in row_entries.iter().enumerate() {
            self.row_ptr[row] = nnz;
            for &(col, value) in row_data {
                self.col_idx.push(col);
                self.values.push(value);
                nnz += 1;
            }
        }
        self.row_ptr[self.nrows] = nnz;

        // Clear assembly map to save memory
        self.assembly_map.clear();
        self.finalized = true;

        Ok(())
    }

    /// Get entry value.
    pub fn get_entry(&self, row: usize, col: usize) -> f64 {
        if !self.finalized {
            // Search in assembly map
            self.assembly_map
                .get(&(row, col))
                .map_or(0.0, |&idx| self.values[idx])
        } else {
            // Search in CSR format
            let start = self.row_ptr[row];
            let end = self.row_ptr[row + 1];

            for i in start..end {
                if self.col_idx[i] == col {
                    return self.values[i];
                }
            }
            0.0
        }
    }

    /// Matrix-vector multiplication: y = A * x.
    pub fn multiply_vector(&self, x: &DVector<f64>) -> FeaResult<DVector<f64>> {
        if !self.finalized {
            return Err(AssemblyError::SparseAssemblyFailed {
                reason: "Matrix must be finalized before multiplication".to_string(),
            }
            .into());
        }

        if x.len() != self.ncols {
            return Err(AssemblyError::MatrixDimensionMismatch {
                expected_rows: self.ncols,
                expected_cols: 1,
                actual_rows: x.len(),
                actual_cols: 1,
            }
            .into());
        }

        let mut y = DVector::zeros(self.nrows);

        for row in 0..self.nrows {
            let start = self.row_ptr[row];
            let end = self.row_ptr[row + 1];

            let mut sum = 0.0;
            for i in start..end {
                sum += self.values[i] * x[self.col_idx[i]];
            }
            y[row] = sum;
        }

        Ok(y)
    }

    /// Extract submatrix for given row and column indices.
    pub fn extract_submatrix(
        &self,
        row_indices: &[usize],
        col_indices: &[usize],
    ) -> FeaResult<Self> {
        if !self.finalized {
            return Err(AssemblyError::SparseAssemblyFailed {
                reason: "Matrix must be finalized before extraction".to_string(),
            }
            .into());
        }

        let sub_nrows = row_indices.len();
        let sub_ncols = col_indices.len();
        let mut sub_matrix = Self::new(sub_nrows, sub_ncols);

        // Create column index mapping
        let mut col_map = HashMap::new();
        for (new_col, &old_col) in col_indices.iter().enumerate() {
            col_map.insert(old_col, new_col);
        }

        for (new_row, &old_row) in row_indices.iter().enumerate() {
            let start = self.row_ptr[old_row];
            let end = self.row_ptr[old_row + 1];

            for i in start..end {
                let old_col = self.col_idx[i];
                if let Some(&new_col) = col_map.get(&old_col) {
                    let value = self.values[i];
                    sub_matrix.add_entry(new_row, new_col, value)?;
                }
            }
        }

        sub_matrix.finalize()?;
        Ok(sub_matrix)
    }

    /// Convert to dense matrix (for debugging/small matrices).
    pub fn to_dense(&self) -> nalgebra::DMatrix<f64> {
        let mut dense = nalgebra::DMatrix::zeros(self.nrows, self.ncols);

        if self.finalized {
            for row in 0..self.nrows {
                let start = self.row_ptr[row];
                let end = self.row_ptr[row + 1];

                for i in start..end {
                    dense[(row, self.col_idx[i])] = self.values[i];
                }
            }
        } else {
            for (&(row, col), &idx) in &self.assembly_map {
                dense[(row, col)] = self.values[idx];
            }
        }

        dense
    }

    /// Get matrix structure (row pointers and column indices).
    pub fn structure(&self) -> (&[usize], &[usize]) {
        (&self.row_ptr, &self.col_idx)
    }

    /// Get matrix values.
    pub fn values(&self) -> &[f64] {
        &self.values
    }

    /// Check if matrix is symmetric (structure-wise).
    pub fn is_symmetric_structure(&self) -> bool {
        if !self.finalized || self.nrows != self.ncols {
            return false;
        }

        // Check if (i,j) exists whenever (j,i) exists
        let mut has_entry = vec![vec![false; self.ncols]; self.nrows];

        for row in 0..self.nrows {
            let start = self.row_ptr[row];
            let end = self.row_ptr[row + 1];

            for i in start..end {
                has_entry[row][self.col_idx[i]] = true;
            }
        }

        for row in 0..self.nrows {
            for col in 0..self.ncols {
                if has_entry[row][col] != has_entry[col][row] {
                    return false;
                }
            }
        }

        true
    }

    /// Check if matrix is symmetric (both structure and values).
    pub fn is_symmetric(&self) -> bool {
        if !self.finalized || self.nrows != self.ncols {
            return false;
        }

        const EPSILON: f64 = 1e-10;

        for row in 0..self.nrows {
            for col in row + 1..self.ncols {
                let val_ij = self.get_entry(row, col);
                let val_ji = self.get_entry(col, row);

                if (val_ij - val_ji).abs() > EPSILON {
                    return false;
                }
            }
        }

        true
    }

    /// Compute matrix norm (Frobenius norm).
    pub fn frobenius_norm(&self) -> f64 {
        self.values.iter().map(|&x| x * x).sum::<f64>().sqrt()
    }

    /// Multiply transpose of matrix with vector: A^T * x
    pub fn transpose_multiply_vector(&self, x: &DVector<f64>) -> FeaResult<DVector<f64>> {
        if x.len() != self.nrows {
            return Err(AssemblyError::DimensionMismatch {
                expected: self.nrows,
                actual: x.len(),
            }
            .into());
        }

        let mut result = DVector::zeros(self.ncols);

        for row in 0..self.nrows {
            for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
                let col = self.col_idx[idx];
                let val = self.values[idx];
                result[col] += val * x[row];
            }
        }

        Ok(result)
    }

    /// Extract diagonal elements
    pub fn diagonal(&self) -> FeaResult<DVector<f64>> {
        let n = self.nrows.min(self.ncols);
        let mut diag = DVector::zeros(n);

        for row in 0..n {
            for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
                if self.col_idx[idx] == row {
                    diag[row] = self.values[idx];
                    break;
                }
            }
        }

        Ok(diag)
    }

    /// Forward substitution for lower triangular system L * x = b
    pub fn forward_solve(&self, b: &DVector<f64>) -> FeaResult<DVector<f64>> {
        if b.len() != self.nrows {
            return Err(AssemblyError::DimensionMismatch {
                expected: self.nrows,
                actual: b.len(),
            }
            .into());
        }

        let mut x = DVector::zeros(self.nrows);

        for row in 0..self.nrows {
            let mut sum = b[row];
            let mut diag_val = 0.0;

            for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
                let col = self.col_idx[idx];
                let val = self.values[idx];

                if col < row {
                    sum -= val * x[col];
                } else if col == row {
                    diag_val = val;
                }
            }

            if diag_val.abs() < 1e-12 {
                return Err(AssemblyError::SingularMatrix.into());
            }

            x[row] = sum / diag_val;
        }

        Ok(x)
    }

    /// Backward substitution for upper triangular system U * x = b
    pub fn backward_solve(&self, b: &DVector<f64>) -> FeaResult<DVector<f64>> {
        if b.len() != self.nrows {
            return Err(AssemblyError::DimensionMismatch {
                expected: self.nrows,
                actual: b.len(),
            }
            .into());
        }

        let mut x = DVector::zeros(self.nrows);

        for row in (0..self.nrows).rev() {
            let mut sum = b[row];
            let mut diag_val = 0.0;

            for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
                let col = self.col_idx[idx];
                let val = self.values[idx];

                if col > row {
                    sum -= val * x[col];
                } else if col == row {
                    diag_val = val;
                }
            }

            if diag_val.abs() < 1e-12 {
                return Err(AssemblyError::SingularMatrix.into());
            }

            x[row] = sum / diag_val;
        }

        Ok(x)
    }

    /// Solve linear system A * x = b using LU decomposition
    pub fn solve_vector(&self, b: &DVector<f64>) -> FeaResult<DVector<f64>> {
        // For now, use a simple iterative method (Conjugate Gradient)
        // In production, this should use a proper sparse solver
        if b.len() != self.nrows || self.nrows != self.ncols {
            return Err(AssemblyError::DimensionMismatch {
                expected: self.nrows,
                actual: b.len(),
            }
            .into());
        }

        // Simple Conjugate Gradient implementation for symmetric positive definite matrices
        let mut x = DVector::zeros(self.nrows);
        let mut r = b.clone();
        let mut p = r.clone();
        let tolerance = 1e-9;
        let max_iter = self.nrows * 2;

        for _iter in 0..max_iter {
            let ap = self.multiply_vector(&p)?;
            let r_dot_r = r.dot(&r);

            if r_dot_r < tolerance * tolerance {
                break;
            }

            let alpha = r_dot_r / p.dot(&ap);
            x += alpha * &p;
            r -= alpha * &ap;

            let r_dot_r_new = r.dot(&r);
            let beta = r_dot_r_new / r_dot_r;
            p = &r + beta * p;
        }

        Ok(x)
    }
}

// Also implement Add for &SparseMatrix + &SparseMatrix for convenience
impl std::ops::Add<&SparseMatrix> for &SparseMatrix {
    type Output = SparseMatrix;

    fn add(self, rhs: &SparseMatrix) -> Self::Output {
        assert!(!(self.nrows != rhs.nrows || self.ncols != rhs.ncols), "Matrix dimensions must match for addition");

        let mut result = SparseMatrix::new(self.nrows, self.ncols);

        // Add entries from left matrix
        if self.finalized {
            // Use CSR format for finalized matrix
            for row in 0..self.nrows {
                for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
                    let col = self.col_idx[idx];
                    let value = self.values[idx];
                    result.add_entry(row, col, value).unwrap();
                }
            }
        } else {
            // Use assembly_map for unfinalized matrix
            for ((row, col), &idx) in &self.assembly_map {
                let value = self.values[idx];
                result.add_entry(*row, *col, value).unwrap();
            }
        }

        // Add entries from right matrix
        if rhs.finalized {
            // Use CSR format for finalized matrix
            for row in 0..rhs.nrows {
                for idx in rhs.row_ptr[row]..rhs.row_ptr[row + 1] {
                    let col = rhs.col_idx[idx];
                    let value = rhs.values[idx];
                    result.add_entry(row, col, value).unwrap();
                }
            }
        } else {
            // Use assembly_map for unfinalized matrix
            for ((row, col), &idx) in &rhs.assembly_map {
                let value = rhs.values[idx];
                result.add_entry(*row, *col, value).unwrap();
            }
        }

        // If both originals were finalized, finalize the result
        if self.finalized && rhs.finalized {
            result.finalize().unwrap();
        }

        result
    }
}

/// GPU-accelerated sparse matrix operations.
#[cfg(feature = "cuda")]
use cudarc::driver::safe::{CudaContext as CudaDevice, CudaStream};

pub struct GpuSparseMatrix {
    /// CPU sparse matrix (for initial data only)
    cpu_matrix: SparseMatrix,
    /// CUDA device - REQUIRED
    #[cfg(feature = "cuda")]
    device: std::sync::Arc<CudaDevice>,
    /// CUDA stream for async operations
    #[cfg(feature = "cuda")]
    stream: std::sync::Arc<CudaStream>,
}

impl GpuSparseMatrix {
    /// Create from CPU sparse matrix.
    /// PANICS if GPU is not available (GPU-only requirement)
    pub fn from_cpu_matrix(cpu_matrix: SparseMatrix) -> FeaResult<Self> {
        #[cfg(feature = "cuda")]
        {
            let device =
                CudaDevice::new(0).expect("GPU is REQUIRED for GpuSparseMatrix. No GPU found.");
            let stream = device.default_stream();

            Ok(Self {
                cpu_matrix,
                device,
                stream,
            })
        }
        #[cfg(not(feature = "cuda"))]
        {
            panic!("GPU is REQUIRED. Compile with 'cuda' feature enabled.")
        }
    }

    /// Transfer matrix to GPU.
    pub fn to_gpu(&mut self) -> FeaResult<()> {
        if !self.cpu_matrix.finalized {
            return Err(AssemblyError::SparseAssemblyFailed {
                reason: "Matrix must be finalized before GPU transfer".to_string(),
            }
            .into());
        }

        #[cfg(feature = "cuda")]
        {
            // GPU is always available here - we panic in from_cpu_matrix if not
            // Allocate device memory for CSR format
            let nnz = self.cpu_matrix.values.len();
            let nrows = self.cpu_matrix.nrows;

            // TODO: Implement actual GPU transfer using cuSPARSE
            // This will use the device and stream fields
            tracing::info!(
                "Matrix transfer to GPU pending implementation (nnz: {}, rows: {})",
                nnz,
                nrows
            );
        }
        Ok(())
    }

    /// GPU matrix-vector multiplication.
    pub fn gpu_multiply_vector(&self, _x: &DVector<f64>) -> FeaResult<DVector<f64>> {
        #[cfg(feature = "cuda")]
        {
            // Perform GPU SpMV using cuSPARSE
            // This would use cusparseSpMV() with optimal buffer allocation
            // and handle device-to-host memory transfers efficiently

            // Steps for full implementation:
            // 1. Allocate device vectors for input/output
            // 2. Copy input vector to device
            // 3. Execute cusparseSpMV operation
            // 4. Copy result back to host

            // TODO: Implement actual GPU SpMV
            // For now, return error to maintain GPU-only requirement
            Err(AssemblyError::SparseAssemblyFailed {
                reason: "GPU SpMV implementation pending".to_string(),
            }
            .into())
        }
        #[cfg(not(feature = "cuda"))]
        {
            panic!("GPU is REQUIRED for matrix operations.")
        }
    }

    /// Get CPU matrix reference.
    pub fn cpu_matrix(&self) -> &SparseMatrix {
        &self.cpu_matrix
    }
}

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

    #[test]
    fn test_sparse_matrix_creation() {
        let matrix = SparseMatrix::new(3, 3);
        assert_eq!(matrix.nrows(), 3);
        assert_eq!(matrix.ncols(), 3);
        assert_eq!(matrix.nnz(), 0);
    }

    #[test]
    fn test_sparse_matrix_assembly() {
        let mut matrix = SparseMatrix::new(3, 3);

        matrix.add_entry(0, 0, 1.0).unwrap();
        matrix.add_entry(1, 1, 2.0).unwrap();
        matrix.add_entry(2, 2, 3.0).unwrap();
        matrix.add_entry(0, 1, 0.5).unwrap();

        assert_eq!(matrix.nnz(), 4);

        // Add to existing entry
        matrix.add_entry(0, 0, 1.0).unwrap();
        assert_eq!(matrix.get_entry(0, 0), 2.0);
    }

    #[test]
    fn test_sparse_matrix_finalization() {
        let mut matrix = SparseMatrix::new(3, 3);

        matrix.add_entry(0, 0, 1.0).unwrap();
        matrix.add_entry(1, 1, 2.0).unwrap();
        matrix.add_entry(2, 2, 3.0).unwrap();

        matrix.finalize().unwrap();

        assert_eq!(matrix.get_entry(0, 0), 1.0);
        assert_eq!(matrix.get_entry(1, 1), 2.0);
        assert_eq!(matrix.get_entry(2, 2), 3.0);
    }

    #[test]
    fn test_matrix_vector_multiplication() {
        let mut matrix = SparseMatrix::new(3, 3);

        matrix.add_entry(0, 0, 2.0).unwrap();
        matrix.add_entry(1, 1, 3.0).unwrap();
        matrix.add_entry(2, 2, 4.0).unwrap();
        matrix.finalize().unwrap();

        let x = DVector::from_vec(vec![1.0, 2.0, 3.0]);
        let y = matrix.multiply_vector(&x).unwrap();

        assert_eq!(y[0], 2.0);
        assert_eq!(y[1], 6.0);
        assert_eq!(y[2], 12.0);
    }

    #[test]
    fn test_submatrix_extraction() {
        let mut matrix = SparseMatrix::new(4, 4);

        // Create a 4x4 diagonal matrix
        for i in 0..4 {
            matrix.add_entry(i, i, (i + 1) as f64).unwrap();
        }
        matrix.finalize().unwrap();

        // Extract 2x2 submatrix
        let row_indices = vec![0, 2];
        let col_indices = vec![0, 2];
        let sub_matrix = matrix
            .extract_submatrix(&row_indices, &col_indices)
            .unwrap();

        assert_eq!(sub_matrix.nrows(), 2);
        assert_eq!(sub_matrix.ncols(), 2);
        assert_eq!(sub_matrix.get_entry(0, 0), 1.0);
        assert_eq!(sub_matrix.get_entry(1, 1), 3.0);
    }

    #[test]
    fn test_dirichlet_boundary_condition() {
        let mut matrix = SparseMatrix::new(3, 3);

        // Fill matrix
        for i in 0..3 {
            for j in 0..3 {
                matrix.add_entry(i, j, ((i + 1) * (j + 1)) as f64).unwrap();
            }
        }

        // Apply Dirichlet BC at DOF 1
        matrix.set_dirichlet(1, 5.0).unwrap();

        // Check that row 1 and column 1 are zeroed except diagonal
        assert_eq!(matrix.get_entry(1, 0), 0.0);
        assert_eq!(matrix.get_entry(1, 1), 1.0);
        assert_eq!(matrix.get_entry(1, 2), 0.0);
        assert_eq!(matrix.get_entry(0, 1), 0.0);
        assert_eq!(matrix.get_entry(2, 1), 0.0);
    }

    #[test]
    fn test_dense_conversion() {
        let mut matrix = SparseMatrix::new(2, 2);
        matrix.add_entry(0, 0, 1.0).unwrap();
        matrix.add_entry(0, 1, 2.0).unwrap();
        matrix.add_entry(1, 0, 3.0).unwrap();
        matrix.add_entry(1, 1, 4.0).unwrap();

        let dense = matrix.to_dense();
        assert_eq!(dense[(0, 0)], 1.0);
        assert_eq!(dense[(0, 1)], 2.0);
        assert_eq!(dense[(1, 0)], 3.0);
        assert_eq!(dense[(1, 1)], 4.0);
    }

    #[test]
    fn test_triplet_construction() {
        let triplets = vec![(0, 0, 1.0), (1, 1, 2.0), (0, 1, 0.5)];

        let matrix = SparseMatrix::from_triplets(2, 2, &triplets).unwrap();
        assert_eq!(matrix.get_entry(0, 0), 1.0);
        assert_eq!(matrix.get_entry(1, 1), 2.0);
        assert_eq!(matrix.get_entry(0, 1), 0.5);
    }

    #[test]
    fn test_matrix_properties() {
        let mut matrix = SparseMatrix::new(3, 3);
        matrix.add_entry(0, 0, 1.0).unwrap();
        matrix.add_entry(1, 1, 4.0).unwrap();
        matrix.add_entry(2, 2, 9.0).unwrap();
        matrix.finalize().unwrap();

        assert_eq!(matrix.nnz(), 3);
        assert!((matrix.density() - 3.0 / 9.0).abs() < 1e-12);
        assert!((matrix.frobenius_norm() - (1.0 + 16.0 + 81.0_f64).sqrt()).abs() < 1e-12);
    }
}

// Implement scalar multiplication for SparseMatrix
impl std::ops::Mul<f64> for SparseMatrix {
    type Output = Self;

    fn mul(self, scalar: f64) -> Self::Output {
        let mut result = Self::new(self.nrows, self.ncols);

        if self.finalized {
            // For finalized matrices, use CSR format data
            for row in 0..self.nrows {
                for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
                    let col = self.col_idx[idx];
                    let value = self.values[idx] * scalar;
                    result.add_entry(row, col, value).unwrap();
                }
            }
            result.finalize().unwrap();
        } else {
            // For unfinalized matrices, use assembly_map
            for ((row, col), &idx) in &self.assembly_map {
                let value = self.values[idx] * scalar;
                result.add_entry(*row, *col, value).unwrap();
            }
        }

        result
    }
}

// Also implement for references to avoid unnecessary cloning
impl std::ops::Mul<f64> for &SparseMatrix {
    type Output = SparseMatrix;

    fn mul(self, scalar: f64) -> Self::Output {
        let mut result = SparseMatrix::new(self.nrows, self.ncols);

        if self.finalized {
            // For finalized matrices, use CSR format data
            for row in 0..self.nrows {
                for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
                    let col = self.col_idx[idx];
                    let value = self.values[idx] * scalar;
                    result.add_entry(row, col, value).unwrap();
                }
            }
            result.finalize().unwrap();
        } else {
            // For unfinalized matrices, use assembly_map
            for ((row, col), &idx) in &self.assembly_map {
                let value = self.values[idx] * scalar;
                result.add_entry(*row, *col, value).unwrap();
            }
        }

        result
    }
}

// Implement matrix addition for SparseMatrix
impl std::ops::Add<&Self> for SparseMatrix {
    type Output = Self;

    fn add(self, rhs: &Self) -> Self::Output {
        assert!(!(self.nrows != rhs.nrows || self.ncols != rhs.ncols), "Matrix dimensions must match for addition");

        let mut result = Self::new(self.nrows, self.ncols);

        // Add entries from left matrix
        if self.finalized {
            // Use CSR format for finalized matrix
            for row in 0..self.nrows {
                for idx in self.row_ptr[row]..self.row_ptr[row + 1] {
                    let col = self.col_idx[idx];
                    let value = self.values[idx];
                    result.add_entry(row, col, value).unwrap();
                }
            }
        } else {
            // Use assembly_map for unfinalized matrix
            for ((row, col), &idx) in &self.assembly_map {
                let value = self.values[idx];
                result.add_entry(*row, *col, value).unwrap();
            }
        }

        // Add entries from right matrix
        if rhs.finalized {
            // Use CSR format for finalized matrix
            for row in 0..rhs.nrows {
                for idx in rhs.row_ptr[row]..rhs.row_ptr[row + 1] {
                    let col = rhs.col_idx[idx];
                    let value = rhs.values[idx];
                    result.add_entry(row, col, value).unwrap();
                }
            }
        } else {
            // Use assembly_map for unfinalized matrix
            for ((row, col), &idx) in &rhs.assembly_map {
                let value = rhs.values[idx];
                result.add_entry(*row, *col, value).unwrap();
            }
        }

        // If both originals were finalized, finalize the result
        if self.finalized && rhs.finalized {
            result.finalize().unwrap();
        }

        result
    }
}