matten 0.4.0

//! The primary public container, [`Tensor`].
//!
//! M2 adds the full creation and conversion surface from RFC-004: fill
//! constructors, `from_vec`, `into_vec`, `arange`/`try_arange`, and
//! `try_from_rows`. Arithmetic, reshape, and I/O arrive in later milestones.

use crate::convert::flatten_rectangular;
use crate::error::MattenError;
use crate::shape;
use std::fmt;

/// Maximum element count accepted by `arange` / `try_arange` before the
/// allocation limit fires. Set conservatively for Phase 1.
const ARANGE_MAX_ELEMENTS: usize = 1 << 28; // ~268 million

/// A dense, row-major, owned multidimensional array of `f64`.
///
/// Fields are private. The storage layout is never exposed across the public
/// API. A scalar is shape `[]` with one element.
///
/// # Examples
///
/// ```
/// use matten::Tensor;
///
/// // 2×2 matrix
/// let m = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
/// assert!(m.is_matrix());
///
/// // Fill constructors
/// let z = Tensor::zeros(&[3]);
/// assert_eq!(z.as_slice(), &[0.0, 0.0, 0.0]);
///
/// // Range constructor
/// let r = Tensor::arange(0.0, 5.0, 1.0);
/// assert_eq!(r.len(), 5);
/// ```
#[derive(Clone, PartialEq)]
pub struct Tensor {
    pub(crate) data: Vec<f64>,
    pub(crate) shape: Vec<usize>,
}

// `is_empty` is absent for 0.1.0: zero-sized dims are rejected and a scalar
// has `len() == 1`, so it would always be false. Deferred to a future
// zero-sized-tensor RFC.
#[allow(clippy::len_without_is_empty)]
impl Tensor {
    // ------------------------------------------------------------------ //
    // Core constructors (also implemented in M1; repeated here for clarity)
    // ------------------------------------------------------------------ //

    /// Creates a tensor from row-major `data` and `shape`.
    ///
    /// Panics with an actionable message on any shape violation or data-length
    /// mismatch. Use [`try_new`](Tensor::try_new) for recoverable construction.
    ///
    /// # Panics
    ///
    /// Panics if the shape is invalid or `data.len()` does not equal the shape
    /// product.
    #[must_use]
    pub fn new(data: Vec<f64>, shape: &[usize]) -> Tensor {
        Self::try_new(data, shape).unwrap_or_else(|e| panic!("{e}"))
    }

    /// Creates a tensor from row-major `data` and `shape`, returning an error
    /// instead of panicking.
    ///
    /// # Errors
    ///
    /// Returns [`MattenError::Shape`] for a length mismatch or invalid shape,
    /// or [`MattenError::Allocation`] on product overflow.
    pub fn try_new(data: Vec<f64>, shape: &[usize]) -> Result<Tensor, MattenError> {
        let expected = shape::validate_shape(shape, "try_new")?;
        if data.len() != expected {
            return Err(MattenError::Shape {
                operation: "try_new",
                message: format!(
                    "data length {} does not match shape {shape:?}, which requires {expected} elements",
                    data.len()
                ),
            });
        }
        Ok(Tensor {
            data,
            shape: shape.to_vec(),
        })
    }

    /// Creates a rank-0 scalar tensor (shape `[]`, length `1`).
    ///
    /// ```
    /// use matten::Tensor;
    /// let s = Tensor::scalar(3.14);
    /// assert!(s.is_scalar());
    /// assert_eq!(s.len(), 1);
    /// ```
    #[must_use]
    pub fn scalar(value: f64) -> Tensor {
        // Shape `[]` is always valid: rank 0, no dims, product 1.
        Tensor {
            data: vec![value],
            shape: Vec::new(),
        }
    }

    // ------------------------------------------------------------------ //
    // Fill constructors
    // ------------------------------------------------------------------ //

    /// Creates a tensor filled with `0.0`.
    ///
    /// # Panics
    ///
    /// Panics on an invalid shape (same rules as [`new`](Tensor::new)).
    ///
    /// ```
    /// use matten::Tensor;
    /// let z = Tensor::zeros(&[2, 3]);
    /// assert_eq!(z.shape(), &[2, 3]);
    /// assert_eq!(z.len(), 6);
    /// assert!(z.as_slice().iter().all(|&v| v == 0.0));
    /// ```
    #[must_use]
    pub fn zeros(shape: &[usize]) -> Tensor {
        let len = shape::validate_shape(shape, "zeros").unwrap_or_else(|e| panic!("{e}"));
        Tensor {
            data: vec![0.0; len],
            shape: shape.to_vec(),
        }
    }

    /// Creates a tensor filled with `1.0`.
    ///
    /// # Panics
    ///
    /// Panics on an invalid shape.
    ///
    /// ```
    /// use matten::Tensor;
    /// let o = Tensor::ones(&[3]);
    /// assert!(o.as_slice().iter().all(|&v| v == 1.0));
    /// ```
    #[must_use]
    pub fn ones(shape: &[usize]) -> Tensor {
        let len = shape::validate_shape(shape, "ones").unwrap_or_else(|e| panic!("{e}"));
        Tensor {
            data: vec![1.0; len],
            shape: shape.to_vec(),
        }
    }

    /// Creates a tensor filled with `value`.
    ///
    /// # Panics
    ///
    /// Panics on an invalid shape.
    ///
    /// ```
    /// use matten::Tensor;
    /// let t = Tensor::full(&[2, 2], 7.0);
    /// assert!(t.as_slice().iter().all(|&v| v == 7.0));
    /// ```
    #[must_use]
    pub fn full(shape: &[usize], value: f64) -> Tensor {
        let len = shape::validate_shape(shape, "full").unwrap_or_else(|e| panic!("{e}"));
        Tensor {
            data: vec![value; len],
            shape: shape.to_vec(),
        }
    }

    /// Creates a 1-D tensor from a flat vector; shape is `[data.len()]`.
    ///
    /// # Panics
    ///
    /// Panics if `data` is empty (zero-sized dimension).
    ///
    /// ```
    /// use matten::Tensor;
    /// let t = Tensor::from_vec(vec![1.0, 2.0, 3.0]);
    /// assert_eq!(t.shape(), &[3]);
    /// ```
    #[must_use]
    pub fn from_vec(data: Vec<f64>) -> Tensor {
        let len = data.len();
        Tensor::new(data, &[len])
    }

    // ------------------------------------------------------------------ //
    // Range constructor
    // ------------------------------------------------------------------ //

    /// Creates a 1-D tensor with values `start, start + step, ...` (half-open,
    /// so `end` is excluded).
    ///
    /// This is a panic-zone convenience. Use [`try_arange`](Tensor::try_arange)
    /// when `start`/`end`/`step` come from user input.
    ///
    /// # Panics
    ///
    /// Panics if `step == 0.0`, any argument is non-finite, or the computed
    /// element count would exceed the allocation limit.
    ///
    /// ```
    /// use matten::Tensor;
    /// let r = Tensor::arange(0.0, 5.0, 1.0);
    /// assert_eq!(r.as_slice(), &[0.0, 1.0, 2.0, 3.0, 4.0]);
    ///
    /// let r2 = Tensor::arange(1.0, 0.0, -0.5);
    /// assert_eq!(r2.as_slice(), &[1.0, 0.5]);
    /// ```
    #[must_use]
    pub fn arange(start: f64, end: f64, step: f64) -> Tensor {
        arange_impl(start, end, step, "arange").unwrap_or_else(|e| panic!("{e}"))
    }

    /// Creates a 1-D tensor with values `start, start + step, ...` (half-open),
    /// returning an error instead of panicking.
    ///
    /// # Errors
    ///
    /// Returns [`MattenError::Shape`] for a zero or non-finite step/bound, or
    /// [`MattenError::Allocation`] if the computed element count exceeds the
    /// allocation limit.
    pub fn try_arange(start: f64, end: f64, step: f64) -> Result<Tensor, MattenError> {
        arange_impl(start, end, step, "try_arange")
    }

    // ------------------------------------------------------------------ //
    // Nested-row construction
    // ------------------------------------------------------------------ //

    /// Creates a 2-D tensor from rectangular row data; shape is `[rows, cols]`.
    ///
    /// This is the recoverable version of [`From<Vec<Vec<f64>>>`].
    ///
    /// # Errors
    ///
    /// Returns [`MattenError::Shape`] on an empty outer vector, any zero-length
    /// row, or ragged rows.
    pub fn try_from_rows(rows: Vec<Vec<f64>>) -> Result<Tensor, MattenError> {
        let (data, shape) = flatten_rectangular(rows, "try_from_rows")?;
        Ok(Tensor { data, shape })
    }

    // ------------------------------------------------------------------ //
    // Observational API (also M1; listed here for doc completeness)
    // ------------------------------------------------------------------ //

    /// The shape as a slice of dimension lengths. Non-allocating.
    #[must_use]
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }

    /// The rank (number of dimensions): `shape().len()`.
    #[must_use]
    pub fn ndim(&self) -> usize {
        self.shape.len()
    }

    /// The logical element count: the product of the shape.
    #[must_use]
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Returns `true` for a rank-0 scalar tensor (shape `[]`).
    #[must_use]
    pub fn is_scalar(&self) -> bool {
        self.ndim() == 0
    }

    /// Returns `true` for a rank-1 tensor.
    #[must_use]
    pub fn is_vector(&self) -> bool {
        self.ndim() == 1
    }

    /// Returns `true` for a rank-2 tensor.
    #[must_use]
    pub fn is_matrix(&self) -> bool {
        self.ndim() == 2
    }

    /// The flat, row-major data as a slice.
    #[must_use]
    pub fn as_slice(&self) -> &[f64] {
        &self.data
    }

    /// Returns an owned copy of the flat, row-major data.
    #[must_use]
    pub fn to_vec(&self) -> Vec<f64> {
        self.data.clone()
    }

    /// Consumes the tensor and returns the flat, row-major data. No copy.
    #[must_use]
    pub fn into_vec(self) -> Vec<f64> {
        self.data
    }
}

impl fmt::Debug for Tensor {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        const MAX: usize = 8;
        write!(f, "Tensor(shape={:?}, data=[", self.shape)?;
        for (i, v) in self.data.iter().take(MAX).enumerate() {
            if i > 0 {
                f.write_str(", ")?;
            }
            write!(f, "{v:?}")?;
        }
        if self.data.len() > MAX {
            write!(f, ", ... ({} more)", self.data.len() - MAX)?;
        }
        f.write_str("])")
    }
}

// ---- arange shared implementation ------------------------------------

/// Shared implementation for both `arange` and `try_arange`.
fn arange_impl(
    start: f64,
    end: f64,
    step: f64,
    operation: &'static str,
) -> Result<Tensor, MattenError> {
    // Argument validation (Result zone, so return Err rather than panic).
    if !start.is_finite() || !end.is_finite() {
        return Err(MattenError::Shape {
            operation,
            message: format!("start and end must be finite (got start={start}, end={end})"),
        });
    }
    if step == 0.0 || !step.is_finite() {
        return Err(MattenError::Shape {
            operation,
            message: format!("step must be a non-zero finite value (got {step})"),
        });
    }

    // Compute element count without allocating. The estimate is conservative:
    // we use integer ceiling so floating-point accumulation never drops a value
    // that should be included.
    let raw_count = ((end - start) / step).ceil();
    let count: usize = if raw_count <= 0.0 {
        0
    } else if raw_count > ARANGE_MAX_ELEMENTS as f64 {
        return Err(MattenError::Allocation {
            requested_elements: raw_count as usize,
            message: format!(
                "arange would produce ~{} elements, exceeding the limit of {}",
                raw_count as usize, ARANGE_MAX_ELEMENTS
            ),
        });
    } else {
        raw_count as usize
    };

    // Build the data vector by accumulation. Recompute to avoid drift near
    // `end`, but exclude values that have genuinely crossed `end`.
    let mut data = Vec::with_capacity(count);
    let mut i: usize = 0;
    loop {
        let v = start + step * i as f64;
        // Half-open: stop when we cross `end` in the direction of `step`.
        if (step > 0.0 && v >= end) || (step < 0.0 && v <= end) {
            break;
        }
        data.push(v);
        i += 1;
        // Guard against the estimate being too small (floating-point edge).
        if i > ARANGE_MAX_ELEMENTS {
            return Err(MattenError::Allocation {
                requested_elements: i,
                message: format!("arange exceeded the element limit of {ARANGE_MAX_ELEMENTS}"),
            });
        }
    }

    let len = data.len();
    // A zero-element range is an edge case we permit as an empty 1-D tensor
    // shape would be [0], which the shape module rejects. Return an error
    // so the user knows to adjust their range arguments.
    if len == 0 {
        return Err(MattenError::Shape {
            operation,
            message: format!("arange(start={start}, end={end}, step={step}) produces no elements"),
        });
    }

    Ok(Tensor {
        data,
        shape: vec![len],
    })
}

impl Tensor {
    // ---- Shape operations (M4 / RFC-007) ------------------------------------

    /// Reshapes the tensor to `new_shape`, returning a new owned tensor.
    ///
    /// The total element count must be unchanged. Data order is preserved
    /// (row-major flat order).
    ///
    /// # Panics
    ///
    /// Panics on element-count mismatch or invalid shape. Use
    /// [`try_reshape`](Tensor::try_reshape) for recoverable construction.
    ///
    /// ```
    /// use matten::Tensor;
    /// let t = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
    /// let flat = t.reshape(&[4]);
    /// assert_eq!(flat.shape(), &[4]);
    /// assert_eq!(flat.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
    /// ```
    #[must_use]
    pub fn reshape(&self, new_shape: &[usize]) -> Tensor {
        crate::reshape::try_reshape_impl(self, new_shape).unwrap_or_else(|e| panic!("{e}"))
    }

    /// Reshapes the tensor, returning an error instead of panicking.
    ///
    /// # Errors
    ///
    /// Returns [`MattenError::Shape`] on element-count mismatch or invalid shape.
    pub fn try_reshape(&self, new_shape: &[usize]) -> Result<Tensor, MattenError> {
        crate::reshape::try_reshape_impl(self, new_shape)
    }

    /// Flattens the tensor to a 1-D tensor, preserving row-major order.
    ///
    /// A scalar (shape `[]`) is returned as shape `[1]`.
    ///
    /// ```
    /// use matten::Tensor;
    /// let t = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
    /// let flat = t.flatten();
    /// assert_eq!(flat.shape(), &[4]);
    /// ```
    #[must_use]
    pub fn flatten(&self) -> Tensor {
        let len = self.data.len();
        Tensor {
            data: self.data.clone(),
            shape: vec![len],
        }
    }

    /// Transposes the tensor by reversing the axis order.
    ///
    /// For a rank-2 tensor this swaps rows and columns. For rank > 2 the axis
    /// order is reversed: `[d0, d1, d2] → [d2, d1, d0]`.
    ///
    /// # Panics
    ///
    /// Panics for a rank-0 scalar (no axes to transpose).
    ///
    /// ```
    /// use matten::Tensor;
    /// let m = Tensor::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
    /// let mt = m.transpose();
    /// assert_eq!(mt.shape(), &[3, 2]);
    /// assert_eq!(mt.as_slice(), &[1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
    /// ```
    #[must_use]
    pub fn transpose(&self) -> Tensor {
        let ndim = self.ndim();
        if ndim == 0 {
            panic!("matten shape error in transpose: cannot transpose a scalar (rank 0)");
        }
        let perm: Vec<usize> = (0..ndim).rev().collect();
        crate::reshape::permute_axes(self, &perm)
    }

    /// Alias for [`transpose`](Tensor::transpose).
    #[must_use]
    pub fn t(&self) -> Tensor {
        self.transpose()
    }

    /// Returns a new tensor with `axis1` and `axis2` swapped.
    ///
    /// # Panics
    ///
    /// Panics if either axis is out of range.
    ///
    /// ```
    /// use matten::Tensor;
    /// let t = Tensor::new((1..=24).map(|x| x as f64).collect(), &[2, 3, 4]);
    /// let s = t.swap_axes(0, 2);
    /// assert_eq!(s.shape(), &[4, 3, 2]);
    /// ```
    #[must_use]
    pub fn swap_axes(&self, axis1: usize, axis2: usize) -> Tensor {
        crate::reshape::validate_axes(axis1, axis2, self.ndim(), "swap_axes")
            .unwrap_or_else(|e| panic!("{e}"));
        let mut perm: Vec<usize> = (0..self.ndim()).collect();
        perm.swap(axis1, axis2);
        crate::reshape::permute_axes(self, &perm)
    }

    /// Returns the element at the multidimensional `coord`, or `None` if the
    /// coordinate rank doesn't match or any component is out of bounds.
    ///
    /// ```
    /// use matten::Tensor;
    /// let t = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
    /// assert_eq!(t.get(&[0, 1]), Some(2.0));
    /// assert_eq!(t.get(&[5, 0]), None);
    /// ```
    pub fn get(&self, coord: &[usize]) -> Option<f64> {
        let flat = crate::shape::coord_to_flat(coord, &self.shape)?;
        self.data.get(flat).copied()
    }

    // ---- Slicing (M4 / RFC-008) ---------------------------------------------

    /// Starts a slice builder for this tensor. The builder is the canonical
    /// slicing API; [`slice_str`](Tensor::slice_str) is a convenience wrapper.
    ///
    /// ```
    /// use matten::Tensor;
    /// let t = Tensor::new(vec![1.0,2.0,3.0,4.0,5.0,6.0], &[2, 3]);
    /// let row = t.slice().index(0).all().build().unwrap();
    /// assert_eq!(row.as_slice(), &[1.0, 2.0, 3.0]);
    /// ```
    pub fn slice(&self) -> crate::slice::SliceBuilder<'_> {
        crate::slice::SliceBuilder::new(self)
    }

    /// Slices this tensor using a NumPy-like string specification.
    ///
    /// This is a convenience wrapper over the builder API. It always returns
    /// `Result` and never panics on malformed input.
    ///
    /// # Errors
    ///
    /// Returns [`MattenError::Slice`] for any parse or bounds error.
    ///
    /// ```
    /// use matten::Tensor;
    /// let t = Tensor::new(vec![1.0,2.0,3.0,4.0,5.0,6.0], &[2, 3]);
    /// let top = t.slice_str("0, :").unwrap();
    /// assert_eq!(top.as_slice(), &[1.0, 2.0, 3.0]);
    /// ```
    pub fn slice_str(&self, spec: &str) -> Result<Tensor, MattenError> {
        let specs = crate::slice::parse_slice_str(spec)?;
        crate::slice::execute_slice(self, &specs, "slice_str")
    }
}