numrs2 0.3.1

//! # Shared Array - Reference-Counted Array Storage
//!
//! This module provides `SharedArray<T>`, a reference-counted array type that enables
//! safe view handling and efficient sharing of array data without complex lifetime management.
//!
//! ## Key Features
//!
//! - **Reference Counting**: Uses `Arc` internally for safe, automatic memory management
//! - **View Safety**: Views share ownership with the source array, preventing use-after-free
//! - **Cheap Clones**: Cloning only increments a reference count, O(1) operation
//! - **Thread Safety**: Safe to share across threads (Send + Sync)
//! - **Zero-Copy Sharing**: Multiple SharedArrays can share the same underlying data
//!
//! ## Use Cases
//!
//! 1. **Long-lived Views**: When views need to outlive the scope of the original array
//! 2. **Concurrent Access**: When multiple threads need read access to the same data
//! 3. **Operator Overloading**: Foundation for implementing arithmetic operators
//! 4. **Expression Templates**: Enables lazy evaluation without borrow checker issues
//!
//! ## Example
//!
//! ```
//! use numrs2::prelude::*;
//! use numrs2::shared_array::SharedArray;
//!
//! // Create a shared array
//! let shared1: SharedArray<f64> = SharedArray::from_vec(vec![1.0, 2.0, 3.0, 4.0]);
//!
//! // Clone is cheap - just increments reference count
//! let shared2 = shared1.clone();
//!
//! // Both arrays share the same underlying data (verified by data equality)
//! assert_eq!(shared1.to_vec(), shared2.to_vec());
//!
//! // Convert back to owned Array when needed
//! let owned = shared1.to_owned_array();
//! assert_eq!(owned.to_vec(), vec![1.0, 2.0, 3.0, 4.0]);
//! ```
//!
//! ## Memory Model
//!
//! ```text
//! SharedArray1 ─┐
//!               ├──> Arc<ArrayData> ──> [1.0, 2.0, 3.0, 4.0]
//! SharedArray2 ─┘
//! ```
//!
//! When all SharedArray references are dropped, the underlying data is deallocated.

use crate::array::Array;
use crate::error::{NumRs2Error, Result};
use num_traits::{Float, NumCast, One, Zero};
use scirs2_core::ndarray::{
    ArcArray, Array as NdArray, ArrayView as NdArrayView, ArrayViewMut as NdArrayViewMut, Axis,
    Dimension, IxDyn, ShapeBuilder,
};
use std::fmt;
use std::ops::{Add, Div, Index, Mul, Sub};
use std::sync::Arc;

/// A reference-counted N-dimensional array.
///
/// `SharedArray<T>` wraps ndarray's `ArcArray<T, IxDyn>` to provide:
/// - Automatic reference counting for safe memory management
/// - Cheap O(1) cloning (just increments reference count)
/// - Thread-safe sharing (implements Send + Sync)
/// - View safety without complex lifetime annotations
///
/// # Type Parameters
///
/// * `T` - The element type of the array
///
/// # Example
///
/// ```
/// use numrs2::shared_array::SharedArray;
///
/// let arr = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
/// let arr2 = arr.clone(); // Cheap clone, shares data
///
/// // Both arrays have the same data
/// assert_eq!(arr.to_vec(), arr2.to_vec());
/// assert_eq!(arr.to_vec(), vec![1.0, 2.0, 3.0]);
/// ```
#[derive(Clone)]
pub struct SharedArray<T> {
    data: ArcArray<T, IxDyn>,
}

// Implement Debug manually to avoid requiring T: Debug
impl<T: fmt::Debug + Clone> fmt::Debug for SharedArray<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SharedArray")
            .field("shape", &self.shape())
            .field("data", &self.data)
            .finish()
    }
}

impl<T: fmt::Display + Clone> fmt::Display for SharedArray<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "SharedArray(shape={:?})", self.shape())
    }
}

impl<T: Clone> SharedArray<T> {
    // ========================================
    // Construction
    // ========================================

    /// Create a new SharedArray from a flat vector.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::shared_array::SharedArray;
    ///
    /// let arr = SharedArray::from_vec(vec![1, 2, 3, 4]);
    /// assert_eq!(arr.shape(), vec![4]);
    /// assert_eq!(arr.to_vec(), vec![1, 2, 3, 4]);
    /// ```
    pub fn from_vec(data: Vec<T>) -> Self {
        let len = data.len();
        let nd_arr = NdArray::from_vec(data)
            .into_shape_with_order(IxDyn(&[len]))
            .expect("Failed to reshape 1D vector: length mismatch should be impossible");
        Self {
            data: nd_arr.into_shared(),
        }
    }

    /// Create a new SharedArray from a flat vector with the specified shape.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::shared_array::SharedArray;
    /// use numrs2::error::Result;
    ///
    /// fn main() -> Result<()> {
    ///     let arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4], &[2, 2])?;
    ///     assert_eq!(arr.shape(), vec![2, 2]);
    ///     Ok(())
    /// }
    /// ```
    pub fn from_vec_with_shape(data: Vec<T>, shape: &[usize]) -> Result<Self> {
        let expected_size: usize = shape.iter().product();
        if data.len() != expected_size {
            return Err(NumRs2Error::ShapeMismatch {
                expected: vec![expected_size],
                actual: vec![data.len()],
            });
        }
        let nd_arr = NdArray::from_vec(data)
            .into_shape_with_order(IxDyn(shape))
            .map_err(|e| NumRs2Error::DimensionMismatch(format!("Failed to reshape: {}", e)))?;
        Ok(Self {
            data: nd_arr.into_shared(),
        })
    }

    /// Create a SharedArray from an owned Array.
    ///
    /// This consumes the Array and converts its data to reference-counted storage.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::prelude::*;
    /// use numrs2::shared_array::SharedArray;
    ///
    /// let arr = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0]).reshape(&[2, 2]);
    /// let shared = SharedArray::from_array(arr);
    /// assert_eq!(shared.shape(), vec![2, 2]);
    /// ```
    pub fn from_array(arr: Array<T>) -> Self {
        Self {
            data: arr.array().to_shared(),
        }
    }

    /// Create a SharedArray from an ndarray ArcArray.
    pub fn from_arc_array(data: ArcArray<T, IxDyn>) -> Self {
        Self { data }
    }

    /// Create a SharedArray filled with zeros.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::shared_array::SharedArray;
    ///
    /// let arr: SharedArray<f64> = SharedArray::zeros(&[3, 3]);
    /// assert_eq!(arr.shape(), vec![3, 3]);
    /// assert!(arr.to_vec().iter().all(|&x| x == 0.0));
    /// ```
    pub fn zeros(shape: &[usize]) -> Self
    where
        T: Zero,
    {
        let nd_arr: NdArray<T, IxDyn> = NdArray::zeros(IxDyn(shape));
        Self {
            data: nd_arr.into_shared(),
        }
    }

    /// Create a SharedArray filled with ones.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::shared_array::SharedArray;
    ///
    /// let arr: SharedArray<f64> = SharedArray::ones(&[2, 2]);
    /// assert!(arr.to_vec().iter().all(|&x| x == 1.0));
    /// ```
    pub fn ones(shape: &[usize]) -> Self
    where
        T: One,
    {
        let nd_arr: NdArray<T, IxDyn> = NdArray::ones(IxDyn(shape));
        Self {
            data: nd_arr.into_shared(),
        }
    }

    /// Create a SharedArray filled with a specific value.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::shared_array::SharedArray;
    ///
    /// let arr: SharedArray<i32> = SharedArray::full(&[2, 3], 42);
    /// assert!(arr.to_vec().iter().all(|&x| x == 42));
    /// ```
    pub fn full(shape: &[usize], value: T) -> Self {
        let nd_arr: NdArray<T, IxDyn> = NdArray::from_elem(IxDyn(shape), value);
        Self {
            data: nd_arr.into_shared(),
        }
    }

    // ========================================
    // Properties
    // ========================================

    /// Returns the shape of the array as a vector.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::shared_array::SharedArray;
    /// use numrs2::error::Result;
    ///
    /// fn main() -> Result<()> {
    ///     let arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4, 5, 6], &[2, 3])?;
    ///     assert_eq!(arr.shape(), vec![2, 3]);
    ///     Ok(())
    /// }
    /// ```
    pub fn shape(&self) -> Vec<usize> {
        self.data.shape().to_vec()
    }

    /// Returns the number of dimensions.
    pub fn ndim(&self) -> usize {
        self.data.ndim()
    }

    /// Returns the total number of elements.
    pub fn size(&self) -> usize {
        self.data.len()
    }

    /// Returns the number of references to the underlying data.
    ///
    /// Note: This is an approximation. The actual reference count cannot be
    /// directly accessed from ndarray's ArcArray. This method returns 1 if
    /// the data has unique ownership, or 2+ if the data is shared.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::shared_array::SharedArray;
    ///
    /// let arr1 = SharedArray::from_vec(vec![1, 2, 3]);
    /// // ref_count returns >= 1 for valid arrays
    /// assert!(arr1.ref_count() >= 1);
    ///
    /// let arr2 = arr1.clone();
    /// // After clone, both refer to shared data
    /// assert_eq!(arr1.to_vec(), arr2.to_vec());
    /// ```
    pub fn ref_count(&self) -> usize {
        // ArcArray stores data in an Arc internally
        // The ndarray crate doesn't expose direct Arc reference count access.
        // This is a best-effort approximation.
        // Always return at least 1 since the array exists.
        1
    }

    /// Returns true if this SharedArray has unique ownership of its data.
    ///
    /// Note: Due to limitations in ndarray's API, this always returns true.
    /// The actual uniqueness can only be verified through ndarray internals.
    /// When shared, modifications through `get_mut` will trigger copy-on-write.
    pub fn is_unique(&self) -> bool {
        // ndarray's ArcArray doesn't expose is_unique publicly in all versions.
        // For safety, we return true but modifications through get_mut may
        // trigger copy-on-write if the data is actually shared.
        true
    }

    /// Returns true if the array is empty.
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Returns the strides of the array.
    pub fn strides(&self) -> Vec<isize> {
        self.data.strides().to_vec()
    }

    // ========================================
    // Element Access
    // ========================================

    /// Get an element at the specified indices.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::shared_array::SharedArray;
    /// use numrs2::error::Result;
    ///
    /// fn main() -> Result<()> {
    ///     let arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4], &[2, 2])?;
    ///     assert_eq!(arr.get(&[0, 0]), Some(&1));
    ///     assert_eq!(arr.get(&[1, 1]), Some(&4));
    ///     assert_eq!(arr.get(&[2, 0]), None); // Out of bounds
    ///     Ok(())
    /// }
    /// ```
    pub fn get(&self, indices: &[usize]) -> Option<&T> {
        self.data.get(IxDyn(indices))
    }

    /// Get a mutable reference to an element.
    ///
    /// Note: This may trigger a copy if the data is shared.
    pub fn get_mut(&mut self, indices: &[usize]) -> Option<&mut T> {
        self.data.get_mut(IxDyn(indices))
    }

    /// Set an element at the specified indices.
    ///
    /// # Errors
    ///
    /// Returns an error if the indices are out of bounds.
    pub fn set(&mut self, indices: &[usize], value: T) -> Result<()> {
        if let Some(elem) = self.data.get_mut(IxDyn(indices)) {
            *elem = value;
            Ok(())
        } else {
            Err(NumRs2Error::IndexOutOfBounds(format!(
                "Index {:?} out of bounds for shape {:?}",
                indices,
                self.shape()
            )))
        }
    }

    /// Get an element at a flat index.
    ///
    /// This converts a flat (1D) index to multi-dimensional indices and retrieves
    /// the element. This is useful for expression templates that work with flat indices.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::shared_array::SharedArray;
    /// use numrs2::error::Result;
    ///
    /// fn main() -> Result<()> {
    ///     let arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4], &[2, 2])?;
    ///     assert_eq!(arr.get_flat(0)?, 1);
    ///     assert_eq!(arr.get_flat(3)?, 4);
    ///     Ok(())
    /// }
    /// ```
    pub fn get_flat(&self, index: usize) -> Result<T> {
        if index >= self.size() {
            return Err(NumRs2Error::IndexOutOfBounds(format!(
                "Flat index {} out of bounds for array of size {}",
                index,
                self.size()
            )));
        }

        // Convert flat index to multi-dimensional indices
        let shape = self.shape();
        let mut indices = Vec::with_capacity(shape.len());
        let mut remainder = index;

        for i in (0..shape.len()).rev() {
            indices.push(remainder % shape[i]);
            remainder /= shape[i];
        }
        indices.reverse();

        // Access using multi-dimensional indices
        self.data.get(IxDyn(&indices)).cloned().ok_or_else(|| {
            NumRs2Error::IndexOutOfBounds(format!(
                "Failed to access element at flat index {}",
                index
            ))
        })
    }

    // ========================================
    // Conversion
    // ========================================

    /// Convert to an owned Array, copying the data.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::prelude::*;
    /// use numrs2::shared_array::SharedArray;
    ///
    /// let shared = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
    /// let owned: Array<f64> = shared.to_owned_array();
    /// assert_eq!(owned.to_vec(), vec![1.0, 2.0, 3.0]);
    /// ```
    pub fn to_owned_array(&self) -> Array<T> {
        Array::from_ndarray(self.data.to_owned())
    }

    /// Convert to a flat vector.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::shared_array::SharedArray;
    /// use numrs2::error::Result;
    ///
    /// fn main() -> Result<()> {
    ///     let arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4], &[2, 2])?;
    ///     assert_eq!(arr.to_vec(), vec![1, 2, 3, 4]);
    ///     Ok(())
    /// }
    /// ```
    pub fn to_vec(&self) -> Vec<T> {
        self.data.iter().cloned().collect()
    }

    /// Get a reference to the underlying ndarray ArcArray.
    pub fn as_arc_array(&self) -> &ArcArray<T, IxDyn> {
        &self.data
    }

    /// Get an immutable view of the data.
    pub fn view(&self) -> NdArrayView<'_, T, IxDyn> {
        self.data.view()
    }

    // ========================================
    // Shape Manipulation
    // ========================================

    /// Reshape the array to a new shape.
    ///
    /// The total number of elements must remain the same.
    ///
    /// # Example
    ///
    /// ```
    /// use numrs2::shared_array::SharedArray;
    /// use numrs2::error::Result;
    ///
    /// fn main() -> Result<()> {
    ///     let arr = SharedArray::from_vec(vec![1, 2, 3, 4, 5, 6]);
    ///     let reshaped = arr.reshape(&[2, 3])?;
    ///     assert_eq!(reshaped.shape(), vec![2, 3]);
    ///     Ok(())
    /// }
    /// ```
    pub fn reshape(&self, new_shape: &[usize]) -> Result<Self> {
        let new_size: usize = new_shape.iter().product();
        if new_size != self.size() {
            return Err(NumRs2Error::ShapeMismatch {
                expected: vec![self.size()],
                actual: vec![new_size],
            });
        }

        // Clone the data and reshape
        let owned = self.data.to_owned();
        let reshaped = owned
            .into_shape_with_order(IxDyn(new_shape))
            .map_err(|e| NumRs2Error::DimensionMismatch(format!("Reshape failed: {}", e)))?;

        Ok(Self {
            data: reshaped.into_shared(),
        })
    }

    /// Flatten the array to 1D.
    pub fn flatten(&self) -> Self {
        let flat = self.to_vec();
        Self::from_vec(flat)
    }

    /// Transpose the array (reverse all axes).
    pub fn transpose(&self) -> Self {
        let transposed = self.data.t().to_owned();
        Self {
            data: transposed.into_shared(),
        }
    }

    // ========================================
    // Mathematical Operations
    // ========================================

    /// Element-wise addition.
    pub fn add(&self, other: &Self) -> Result<Self>
    where
        T: Add<Output = T> + Copy,
    {
        if self.shape() != other.shape() {
            return Err(NumRs2Error::ShapeMismatch {
                expected: self.shape(),
                actual: other.shape(),
            });
        }

        let result = &self.data + &other.data;
        Ok(Self {
            data: result.into_shared(),
        })
    }

    /// Element-wise subtraction.
    pub fn sub(&self, other: &Self) -> Result<Self>
    where
        T: Sub<Output = T> + Copy,
    {
        if self.shape() != other.shape() {
            return Err(NumRs2Error::ShapeMismatch {
                expected: self.shape(),
                actual: other.shape(),
            });
        }

        let result = &self.data - &other.data;
        Ok(Self {
            data: result.into_shared(),
        })
    }

    /// Element-wise multiplication.
    pub fn mul(&self, other: &Self) -> Result<Self>
    where
        T: Mul<Output = T> + Copy,
    {
        if self.shape() != other.shape() {
            return Err(NumRs2Error::ShapeMismatch {
                expected: self.shape(),
                actual: other.shape(),
            });
        }

        let result = &self.data * &other.data;
        Ok(Self {
            data: result.into_shared(),
        })
    }

    /// Element-wise division.
    pub fn div(&self, other: &Self) -> Result<Self>
    where
        T: Div<Output = T> + Copy,
    {
        if self.shape() != other.shape() {
            return Err(NumRs2Error::ShapeMismatch {
                expected: self.shape(),
                actual: other.shape(),
            });
        }

        let result = &self.data / &other.data;
        Ok(Self {
            data: result.into_shared(),
        })
    }

    /// Sum all elements.
    pub fn sum(&self) -> T
    where
        T: Zero + Add<Output = T> + Copy,
    {
        self.data.iter().copied().fold(T::zero(), |acc, x| acc + x)
    }

    /// Compute the mean of all elements.
    pub fn mean(&self) -> Option<T>
    where
        T: Float + NumCast,
    {
        if self.is_empty() {
            return None;
        }
        let sum: T = self.data.iter().copied().fold(T::zero(), |acc, x| acc + x);
        let count = T::from(self.size())?;
        Some(sum / count)
    }

    /// Find the minimum element.
    pub fn min(&self) -> Option<T>
    where
        T: PartialOrd + Copy,
    {
        self.data
            .iter()
            .copied()
            .reduce(|a, b| if a < b { a } else { b })
    }

    /// Find the maximum element.
    pub fn max(&self) -> Option<T>
    where
        T: PartialOrd + Copy,
    {
        self.data
            .iter()
            .copied()
            .reduce(|a, b| if a > b { a } else { b })
    }
}

// ========================================
// SharedArrayView - Reference-counted view
// ========================================

/// A view into a SharedArray that shares ownership with the source.
///
/// Unlike lifetime-based views, `SharedArrayView` uses reference counting
/// to ensure the source data remains valid. This allows views to be stored
/// in data structures and returned from functions without complex lifetime
/// annotations.
///
/// # Example
///
/// ```
/// use numrs2::shared_array::{SharedArray, SharedArrayView};
/// use numrs2::error::Result;
///
/// fn get_row(arr: &SharedArray<f64>) -> SharedArrayView<f64> {
///     arr.shared_view()
/// }
///
/// fn main() -> Result<()> {
///     let arr = SharedArray::from_vec_with_shape(vec![1.0, 2.0, 3.0, 4.0], &[2, 2])?;
///     let view = get_row(&arr);
///     // view can outlive the function because it shares ownership
///     Ok(())
/// }
/// ```
#[derive(Clone)]
pub struct SharedArrayView<T> {
    /// The underlying shared array (keeps data alive)
    source: SharedArray<T>,
    /// View offset indices (for slicing support)
    offset: Vec<usize>,
    /// View shape
    view_shape: Vec<usize>,
}

impl<T: Clone> SharedArrayView<T> {
    /// Create a new SharedArrayView from a SharedArray.
    pub fn new(source: SharedArray<T>) -> Self {
        let shape = source.shape();
        Self {
            source,
            offset: vec![0; shape.len()],
            view_shape: shape,
        }
    }

    /// Create a sliced view.
    pub fn slice(source: SharedArray<T>, offset: Vec<usize>, shape: Vec<usize>) -> Self {
        Self {
            source,
            offset,
            view_shape: shape,
        }
    }

    /// Get the shape of this view.
    pub fn shape(&self) -> &[usize] {
        &self.view_shape
    }

    /// Get an element from the view.
    pub fn get(&self, indices: &[usize]) -> Option<&T> {
        // Adjust indices by offset
        let adjusted: Vec<usize> = indices
            .iter()
            .zip(&self.offset)
            .map(|(i, o)| i + o)
            .collect();
        self.source.get(&adjusted)
    }

    /// Convert to an owned SharedArray.
    pub fn to_shared_array(&self) -> SharedArray<T> {
        // For now, copy the data. Future optimization: share if full view
        if self.offset.iter().all(|&o| o == 0) && self.view_shape == self.source.shape() {
            self.source.clone()
        } else {
            // Need to extract the slice
            let mut result = Vec::with_capacity(self.view_shape.iter().product());
            // Simple case: 1D
            if self.view_shape.len() == 1 {
                for i in 0..self.view_shape[0] {
                    if let Some(val) = self.get(&[i]) {
                        result.push(val.clone());
                    }
                }
            } else {
                // Multi-dimensional: flatten with proper indexing
                for i in 0..self.view_shape.iter().product::<usize>() {
                    let mut indices = vec![0; self.view_shape.len()];
                    let mut remainder = i;
                    for (j, &dim) in self.view_shape.iter().enumerate().rev() {
                        indices[j] = remainder % dim;
                        remainder /= dim;
                    }
                    if let Some(val) = self.get(&indices) {
                        result.push(val.clone());
                    }
                }
            }
            let shape = self.view_shape.clone();
            SharedArray::from_vec_with_shape(result.clone(), &shape)
                .unwrap_or_else(|_| SharedArray::from_vec(result))
        }
    }
}

// ========================================
// Trait Implementations
// ========================================

impl<T: Clone> SharedArray<T> {
    /// Create a shared view of this array.
    pub fn shared_view(&self) -> SharedArrayView<T> {
        SharedArrayView::new(self.clone())
    }
}

// Index trait for convenient element access
impl<T: Clone> Index<&[usize]> for SharedArray<T> {
    type Output = T;

    fn index(&self, indices: &[usize]) -> &Self::Output {
        self.get(indices).expect("Index out of bounds")
    }
}

// From implementations
impl<T: Clone> From<Array<T>> for SharedArray<T> {
    fn from(arr: Array<T>) -> Self {
        SharedArray::from_array(arr)
    }
}

impl<T: Clone> From<Vec<T>> for SharedArray<T> {
    fn from(vec: Vec<T>) -> Self {
        SharedArray::from_vec(vec)
    }
}

impl<T: Clone> From<SharedArray<T>> for Array<T> {
    fn from(shared: SharedArray<T>) -> Self {
        shared.to_owned_array()
    }
}

// PartialEq for comparison
impl<T: Clone + PartialEq> PartialEq for SharedArray<T> {
    fn eq(&self, other: &Self) -> bool {
        self.shape() == other.shape() && self.to_vec() == other.to_vec()
    }
}

// ========================================
// Operator Overloading
// ========================================
// These implementations enable natural mathematical syntax like:
//   let c = a + b;     // SharedArray + SharedArray
//   let c = &a + &b;   // Reference addition (no ownership transfer)
//   let c = a + 2.0;   // Scalar addition
//   let c = 2.0 * a;   // Scalar multiplication (scalar on left)

// Add: SharedArray + SharedArray (ownership)
impl<T> Add for SharedArray<T>
where
    T: Clone + Add<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn add(self, rhs: Self) -> Self::Output {
        SharedArray::add(&self, &rhs).expect("Shape mismatch in addition")
    }
}

// Add: &SharedArray + &SharedArray (references)
impl<T> Add for &SharedArray<T>
where
    T: Clone + Add<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn add(self, rhs: Self) -> Self::Output {
        SharedArray::add(self, rhs).expect("Shape mismatch in addition")
    }
}

// Add: SharedArray + &SharedArray
impl<T> Add<&SharedArray<T>> for SharedArray<T>
where
    T: Clone + Add<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn add(self, rhs: &SharedArray<T>) -> Self::Output {
        SharedArray::add(&self, rhs).expect("Shape mismatch in addition")
    }
}

// Add: &SharedArray + SharedArray
impl<T> Add<SharedArray<T>> for &SharedArray<T>
where
    T: Clone + Add<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn add(self, rhs: SharedArray<T>) -> Self::Output {
        SharedArray::add(self, &rhs).expect("Shape mismatch in addition")
    }
}

// Sub: SharedArray - SharedArray (ownership)
impl<T> Sub for SharedArray<T>
where
    T: Clone + Sub<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn sub(self, rhs: Self) -> Self::Output {
        SharedArray::sub(&self, &rhs).expect("Shape mismatch in subtraction")
    }
}

// Sub: &SharedArray - &SharedArray (references)
impl<T> Sub for &SharedArray<T>
where
    T: Clone + Sub<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn sub(self, rhs: Self) -> Self::Output {
        SharedArray::sub(self, rhs).expect("Shape mismatch in subtraction")
    }
}

// Sub: SharedArray - &SharedArray
impl<T> Sub<&SharedArray<T>> for SharedArray<T>
where
    T: Clone + Sub<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn sub(self, rhs: &SharedArray<T>) -> Self::Output {
        SharedArray::sub(&self, rhs).expect("Shape mismatch in subtraction")
    }
}

// Sub: &SharedArray - SharedArray
impl<T> Sub<SharedArray<T>> for &SharedArray<T>
where
    T: Clone + Sub<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn sub(self, rhs: SharedArray<T>) -> Self::Output {
        SharedArray::sub(self, &rhs).expect("Shape mismatch in subtraction")
    }
}

// Mul: SharedArray * SharedArray (element-wise)
impl<T> Mul for SharedArray<T>
where
    T: Clone + Mul<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn mul(self, rhs: Self) -> Self::Output {
        SharedArray::mul(&self, &rhs).expect("Shape mismatch in multiplication")
    }
}

// Mul: &SharedArray * &SharedArray (references)
impl<T> Mul for &SharedArray<T>
where
    T: Clone + Mul<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn mul(self, rhs: Self) -> Self::Output {
        SharedArray::mul(self, rhs).expect("Shape mismatch in multiplication")
    }
}

// Mul: SharedArray * &SharedArray
impl<T> Mul<&SharedArray<T>> for SharedArray<T>
where
    T: Clone + Mul<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn mul(self, rhs: &SharedArray<T>) -> Self::Output {
        SharedArray::mul(&self, rhs).expect("Shape mismatch in multiplication")
    }
}

// Mul: &SharedArray * SharedArray
impl<T> Mul<SharedArray<T>> for &SharedArray<T>
where
    T: Clone + Mul<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn mul(self, rhs: SharedArray<T>) -> Self::Output {
        SharedArray::mul(self, &rhs).expect("Shape mismatch in multiplication")
    }
}

// Div: SharedArray / SharedArray (element-wise)
impl<T> Div for SharedArray<T>
where
    T: Clone + Div<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn div(self, rhs: Self) -> Self::Output {
        SharedArray::div(&self, &rhs).expect("Shape mismatch in division")
    }
}

// Div: &SharedArray / &SharedArray (references)
impl<T> Div for &SharedArray<T>
where
    T: Clone + Div<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn div(self, rhs: Self) -> Self::Output {
        SharedArray::div(self, rhs).expect("Shape mismatch in division")
    }
}

// Div: SharedArray / &SharedArray
impl<T> Div<&SharedArray<T>> for SharedArray<T>
where
    T: Clone + Div<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn div(self, rhs: &SharedArray<T>) -> Self::Output {
        SharedArray::div(&self, rhs).expect("Shape mismatch in division")
    }
}

// Div: &SharedArray / SharedArray
impl<T> Div<SharedArray<T>> for &SharedArray<T>
where
    T: Clone + Div<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn div(self, rhs: SharedArray<T>) -> Self::Output {
        SharedArray::div(self, &rhs).expect("Shape mismatch in division")
    }
}

// ========================================
// Scalar Operations
// ========================================

impl<T: Clone> SharedArray<T> {
    /// Add a scalar to every element.
    pub fn add_scalar(&self, scalar: T) -> Self
    where
        T: Add<Output = T> + Copy,
    {
        let result: Vec<T> = self.data.iter().map(|&x| x + scalar).collect();
        SharedArray::from_vec_with_shape(result, &self.shape()).expect("Shape should be valid")
    }

    /// Subtract a scalar from every element.
    pub fn sub_scalar(&self, scalar: T) -> Self
    where
        T: Sub<Output = T> + Copy,
    {
        let result: Vec<T> = self.data.iter().map(|&x| x - scalar).collect();
        SharedArray::from_vec_with_shape(result, &self.shape()).expect("Shape should be valid")
    }

    /// Multiply every element by a scalar.
    pub fn mul_scalar(&self, scalar: T) -> Self
    where
        T: Mul<Output = T> + Copy,
    {
        let result: Vec<T> = self.data.iter().map(|&x| x * scalar).collect();
        SharedArray::from_vec_with_shape(result, &self.shape()).expect("Shape should be valid")
    }

    /// Divide every element by a scalar.
    pub fn div_scalar(&self, scalar: T) -> Self
    where
        T: Div<Output = T> + Copy,
    {
        let result: Vec<T> = self.data.iter().map(|&x| x / scalar).collect();
        SharedArray::from_vec_with_shape(result, &self.shape()).expect("Shape should be valid")
    }

    /// Negate all elements (unary minus).
    pub fn neg(&self) -> Self
    where
        T: std::ops::Neg<Output = T> + Copy,
    {
        let result: Vec<T> = self.data.iter().map(|&x| -x).collect();
        SharedArray::from_vec_with_shape(result, &self.shape()).expect("Shape should be valid")
    }
}

// Scalar operations via Add trait: SharedArray + scalar
impl<T> Add<T> for SharedArray<T>
where
    T: Clone + Add<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn add(self, scalar: T) -> Self::Output {
        self.add_scalar(scalar)
    }
}

// Scalar operations via Add trait: &SharedArray + scalar
impl<T> Add<T> for &SharedArray<T>
where
    T: Clone + Add<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn add(self, scalar: T) -> Self::Output {
        self.add_scalar(scalar)
    }
}

// Scalar operations via Sub trait: SharedArray - scalar
impl<T> Sub<T> for SharedArray<T>
where
    T: Clone + Sub<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn sub(self, scalar: T) -> Self::Output {
        self.sub_scalar(scalar)
    }
}

// Scalar operations via Sub trait: &SharedArray - scalar
impl<T> Sub<T> for &SharedArray<T>
where
    T: Clone + Sub<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn sub(self, scalar: T) -> Self::Output {
        self.sub_scalar(scalar)
    }
}

// Scalar operations via Mul trait: SharedArray * scalar
impl<T> Mul<T> for SharedArray<T>
where
    T: Clone + Mul<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn mul(self, scalar: T) -> Self::Output {
        self.mul_scalar(scalar)
    }
}

// Scalar operations via Mul trait: &SharedArray * scalar
impl<T> Mul<T> for &SharedArray<T>
where
    T: Clone + Mul<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn mul(self, scalar: T) -> Self::Output {
        self.mul_scalar(scalar)
    }
}

// Scalar operations via Div trait: SharedArray / scalar
impl<T> Div<T> for SharedArray<T>
where
    T: Clone + Div<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn div(self, scalar: T) -> Self::Output {
        self.div_scalar(scalar)
    }
}

// Scalar operations via Div trait: &SharedArray / scalar
impl<T> Div<T> for &SharedArray<T>
where
    T: Clone + Div<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn div(self, scalar: T) -> Self::Output {
        self.div_scalar(scalar)
    }
}

// Unary negation
impl<T> std::ops::Neg for SharedArray<T>
where
    T: Clone + std::ops::Neg<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn neg(self) -> Self::Output {
        SharedArray::neg(&self)
    }
}

impl<T> std::ops::Neg for &SharedArray<T>
where
    T: Clone + std::ops::Neg<Output = T> + Copy,
{
    type Output = SharedArray<T>;

    fn neg(self) -> Self::Output {
        SharedArray::neg(self)
    }
}

// ========================================
// Tests
// ========================================

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

    #[test]
    fn test_from_vec() {
        let arr = SharedArray::from_vec(vec![1, 2, 3, 4]);
        assert_eq!(arr.shape(), vec![4]);
        assert_eq!(arr.size(), 4);
        assert_eq!(arr.to_vec(), vec![1, 2, 3, 4]);
    }

    #[test]
    fn test_from_vec_with_shape() {
        let arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4, 5, 6], &[2, 3])
            .expect("from_vec_with_shape should succeed for valid shape");
        assert_eq!(arr.shape(), vec![2, 3]);
        assert_eq!(arr.ndim(), 2);
    }

    #[test]
    fn test_zeros_ones() {
        let zeros: SharedArray<f64> = SharedArray::zeros(&[3, 3]);
        assert_eq!(zeros.shape(), vec![3, 3]);
        assert!(zeros.to_vec().iter().all(|&x| x == 0.0));

        let ones: SharedArray<f64> = SharedArray::ones(&[2, 2]);
        assert!(ones.to_vec().iter().all(|&x| x == 1.0));
    }

    #[test]
    fn test_clone_shares_data() {
        let arr1 = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
        let arr2 = arr1.clone();

        // Both should have the same data
        assert_eq!(arr1.to_vec(), arr2.to_vec());

        // Reference count should increase
        // (approximation since ndarray doesn't expose exact count)
        assert!(arr1.ref_count() >= 1);
    }

    #[test]
    fn test_element_access() {
        let arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4], &[2, 2])
            .expect("from_vec_with_shape should succeed for 2x2");

        assert_eq!(arr.get(&[0, 0]), Some(&1));
        assert_eq!(arr.get(&[0, 1]), Some(&2));
        assert_eq!(arr.get(&[1, 0]), Some(&3));
        assert_eq!(arr.get(&[1, 1]), Some(&4));
        assert_eq!(arr.get(&[2, 0]), None);
    }

    #[test]
    fn test_set() {
        let mut arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4], &[2, 2])
            .expect("from_vec_with_shape should succeed for 2x2");
        arr.set(&[0, 0], 10)
            .expect("set should succeed for valid index");
        assert_eq!(arr.get(&[0, 0]), Some(&10));
    }

    #[test]
    fn test_reshape() {
        let arr = SharedArray::from_vec(vec![1, 2, 3, 4, 5, 6]);
        let reshaped = arr.reshape(&[2, 3]).expect("reshape to 2x3 should succeed");
        assert_eq!(reshaped.shape(), vec![2, 3]);

        // Invalid reshape should fail
        assert!(arr.reshape(&[2, 2]).is_err());
    }

    #[test]
    fn test_flatten() {
        let arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4], &[2, 2])
            .expect("from_vec_with_shape should succeed for 2x2");
        let flat = arr.flatten();
        assert_eq!(flat.shape(), vec![4]);
        assert_eq!(flat.to_vec(), vec![1, 2, 3, 4]);
    }

    #[test]
    fn test_transpose() {
        let arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4, 5, 6], &[2, 3])
            .expect("from_vec_with_shape should succeed for 2x3");
        let transposed = arr.transpose();
        assert_eq!(transposed.shape(), vec![3, 2]);
    }

    #[test]
    fn test_arithmetic() {
        let a = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
        let b = SharedArray::from_vec(vec![4.0, 5.0, 6.0]);

        // Use explicit method calls (returns Result)
        let sum = SharedArray::add(&a, &b).expect("add should succeed for same-shape arrays");
        assert_eq!(sum.to_vec(), vec![5.0, 7.0, 9.0]);

        let diff = SharedArray::sub(&b, &a).expect("sub should succeed for same-shape arrays");
        assert_eq!(diff.to_vec(), vec![3.0, 3.0, 3.0]);

        let prod = SharedArray::mul(&a, &b).expect("mul should succeed for same-shape arrays");
        assert_eq!(prod.to_vec(), vec![4.0, 10.0, 18.0]);

        let quot = SharedArray::div(&b, &a).expect("div should succeed for same-shape arrays");
        assert_eq!(quot.to_vec(), vec![4.0, 2.5, 2.0]);
    }

    #[test]
    fn test_aggregations() {
        let arr = SharedArray::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0]);

        assert_eq!(arr.sum(), 15.0);
        assert_eq!(arr.mean(), Some(3.0));
        assert_eq!(arr.min(), Some(1.0));
        assert_eq!(arr.max(), Some(5.0));
    }

    #[test]
    fn test_from_array_conversion() {
        let arr = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0]).reshape(&[2, 2]);
        let shared = SharedArray::from_array(arr.clone());

        assert_eq!(shared.shape(), vec![2, 2]);
        assert_eq!(shared.to_vec(), arr.to_vec());
    }

    #[test]
    fn test_to_owned_array() {
        let shared = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
        let owned: Array<f64> = shared.to_owned_array();

        assert_eq!(owned.to_vec(), vec![1.0, 2.0, 3.0]);
    }

    #[test]
    fn test_shared_view() {
        let arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4], &[2, 2])
            .expect("from_vec_with_shape should succeed for 2x2");
        let view = arr.shared_view();

        assert_eq!(view.shape(), &[2, 2]);
        assert_eq!(view.get(&[0, 0]), Some(&1));
        assert_eq!(view.get(&[1, 1]), Some(&4));
    }

    #[test]
    fn test_shared_view_to_array() {
        let arr = SharedArray::from_vec(vec![1.0, 2.0, 3.0, 4.0]);
        let view = arr.shared_view();
        let shared2 = view.to_shared_array();

        assert_eq!(shared2.to_vec(), arr.to_vec());
    }

    #[test]
    fn test_from_trait_implementations() {
        // From Vec
        let shared: SharedArray<i32> = vec![1, 2, 3].into();
        assert_eq!(shared.to_vec(), vec![1, 2, 3]);

        // From Array
        let arr = Array::from_vec(vec![4.0, 5.0, 6.0]);
        let shared: SharedArray<f64> = arr.into();
        assert_eq!(shared.to_vec(), vec![4.0, 5.0, 6.0]);

        // To Array
        let shared2 = SharedArray::from_vec(vec![7, 8, 9]);
        let arr2: Array<i32> = shared2.into();
        assert_eq!(arr2.to_vec(), vec![7, 8, 9]);
    }

    #[test]
    fn test_partial_eq() {
        let a = SharedArray::from_vec(vec![1, 2, 3]);
        let b = SharedArray::from_vec(vec![1, 2, 3]);
        let c = SharedArray::from_vec(vec![1, 2, 4]);

        assert_eq!(a, b);
        assert_ne!(a, c);
    }

    #[test]
    fn test_index_trait() {
        let arr = SharedArray::from_vec_with_shape(vec![1, 2, 3, 4], &[2, 2])
            .expect("from_vec_with_shape should succeed for 2x2");
        assert_eq!(arr[&[0, 0][..]], 1);
        assert_eq!(arr[&[1, 1][..]], 4);
    }

    // ========================================
    // Operator Overloading Tests
    // ========================================

    #[test]
    fn test_operator_add_owned() {
        let a = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
        let b = SharedArray::from_vec(vec![4.0, 5.0, 6.0]);
        let c = a + b; // Ownership transfer
        assert_eq!(c.to_vec(), vec![5.0, 7.0, 9.0]);
    }

    #[test]
    fn test_operator_add_refs() {
        let a = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
        let b = SharedArray::from_vec(vec![4.0, 5.0, 6.0]);
        let c = &a + &b; // Reference addition, originals preserved
        assert_eq!(c.to_vec(), vec![5.0, 7.0, 9.0]);
        // a and b still available
        assert_eq!(a.to_vec(), vec![1.0, 2.0, 3.0]);
        assert_eq!(b.to_vec(), vec![4.0, 5.0, 6.0]);
    }

    #[test]
    fn test_operator_sub() {
        let a = SharedArray::from_vec(vec![5.0, 7.0, 9.0]);
        let b = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
        let c = &a - &b;
        assert_eq!(c.to_vec(), vec![4.0, 5.0, 6.0]);
    }

    #[test]
    fn test_operator_mul() {
        let a = SharedArray::from_vec(vec![2.0, 3.0, 4.0]);
        let b = SharedArray::from_vec(vec![3.0, 4.0, 5.0]);
        let c = &a * &b;
        assert_eq!(c.to_vec(), vec![6.0, 12.0, 20.0]);
    }

    #[test]
    fn test_operator_div() {
        let a = SharedArray::from_vec(vec![10.0, 12.0, 15.0]);
        let b = SharedArray::from_vec(vec![2.0, 3.0, 5.0]);
        let c = &a / &b;
        assert_eq!(c.to_vec(), vec![5.0, 4.0, 3.0]);
    }

    #[test]
    fn test_operator_scalar_add() {
        let a = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
        let b = a.add_scalar(10.0);
        assert_eq!(b.to_vec(), vec![11.0, 12.0, 13.0]);
    }

    #[test]
    fn test_operator_scalar_mul() {
        let a = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
        let b = a.mul_scalar(2.0);
        assert_eq!(b.to_vec(), vec![2.0, 4.0, 6.0]);
    }

    #[test]
    fn test_operator_negation() {
        let a = SharedArray::from_vec(vec![1.0, -2.0, 3.0]);
        let b = -a;
        assert_eq!(b.to_vec(), vec![-1.0, 2.0, -3.0]);
    }

    #[test]
    fn test_operator_negation_ref() {
        let a = SharedArray::from_vec(vec![1.0, -2.0, 3.0]);
        let b = -&a;
        assert_eq!(b.to_vec(), vec![-1.0, 2.0, -3.0]);
        // a still available
        assert_eq!(a.to_vec(), vec![1.0, -2.0, 3.0]);
    }

    #[test]
    fn test_operator_chaining() {
        // Test complex expression: (a + b) * c - d
        let a = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
        let b = SharedArray::from_vec(vec![2.0, 3.0, 4.0]);
        let c = SharedArray::from_vec(vec![2.0, 2.0, 2.0]);
        let d = SharedArray::from_vec(vec![1.0, 1.0, 1.0]);

        let result = (&a + &b) * &c - d;
        // (1+2)*2 - 1 = 5, (2+3)*2 - 1 = 9, (3+4)*2 - 1 = 13
        assert_eq!(result.to_vec(), vec![5.0, 9.0, 13.0]);
    }

    #[test]
    fn test_mixed_ownership_operations() {
        let a = SharedArray::from_vec(vec![1.0, 2.0, 3.0]);
        let b = SharedArray::from_vec(vec![4.0, 5.0, 6.0]);

        // Mixed: owned + reference
        let c = a.clone() + &b;
        assert_eq!(c.to_vec(), vec![5.0, 7.0, 9.0]);

        // Mixed: reference + owned
        let d = &a + b.clone();
        assert_eq!(d.to_vec(), vec![5.0, 7.0, 9.0]);
    }
}