// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// Tensor.hpp
// Tensor class of arbitrary rank
//
// Features:
//   - Dynamic rank and dynamic shape
//   - Stride-based view semantics (transpose/reshape avoid copying data)
//   - Data sharing via shared_ptr (views safely reference the original storage)
//   - Conversion to and from Matrix<T> and Vector<T>
//   - Type constraints via C++23 concepts

#ifndef SANGI_TENSOR_HPP
#define SANGI_TENSOR_HPP

#include <math/core/common.hpp>
#include <math/core/traits.hpp>

#include <algorithm>
#include <cstddef>
#include <initializer_list>
#include <iosfwd>
#include <memory>
#include <numeric>
#include <sstream>
#include <span>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace sangi {

    // Forward declarations
    template<typename T> class Matrix;
    template<typename T> class Vector;

    // ================================================================
    // Requirements on the element type of Tensor<T>
    // ================================================================

    template<typename T>
    concept TensorScalar = requires(T a, T b) {
        { a + b } -> std::convertible_to<T>;
        { a - b } -> std::convertible_to<T>;
        { a * b } -> std::convertible_to<T>;
        { -a } -> std::convertible_to<T>;
        { T(0) };
    };

    // ================================================================
    // Stride computation utilities (detail)
    // ================================================================

    namespace detail {

        // C-order (row-major) stride computation
        inline std::vector<size_t> computeRowMajorStrides(const std::vector<size_t>& shape) {
            if (shape.empty()) return {};
            std::vector<size_t> strides(shape.size());
            strides.back() = 1;
            for (size_t i = shape.size() - 1; i > 0; --i) {
                strides[i - 1] = strides[i] * shape[i];
            }
            return strides;
        }

        // Compute the total element count from shape
        inline size_t computeTotalSize(const std::vector<size_t>& shape) {
            if (shape.empty()) return 1; // a rank-0 scalar has exactly 1 element
            return std::accumulate(shape.begin(), shape.end(), size_t(1), std::multiplies<size_t>());
        }

        // Determine whether the strides are contiguous (C-order)
        inline bool isContiguousStrides(const std::vector<size_t>& shape, const std::vector<size_t>& strides) {
            if (shape.empty()) return true;
            auto expected = computeRowMajorStrides(shape);
            return strides == expected;
        }

    } // namespace detail

    // ================================================================
    // Tensor<T> class
    // ================================================================

    template<TensorScalar T>
    class Tensor {
    public:
        using value_type = T;
        using size_type = std::size_t;
        using reference = T&;
        using const_reference = const T&;
        using pointer = T*;
        using const_pointer = const T*;

        // ============================================================
        // Constructors
        // ============================================================

        // Default: rank-0 scalar (value = T(0))
        Tensor()
            : data_(std::make_shared<std::vector<T>>(1, T(0)))
            , shape_()
            , strides_()
            , offset_(0) {
        }

        // With explicit shape (zero-initialized)
        explicit Tensor(std::vector<size_t> shape)
            : data_(std::make_shared<std::vector<T>>(detail::computeTotalSize(shape), T(0)))
            , shape_(std::move(shape))
            , strides_(detail::computeRowMajorStrides(shape_))
            , offset_(0) {
        }

        // With explicit shape (initializer_list overload, for overload resolution)
        explicit Tensor(std::initializer_list<size_t> shape)
            : Tensor(std::vector<size_t>(shape)) {
        }

        // shape + initial value
        Tensor(std::vector<size_t> shape, const T& value)
            : data_(std::make_shared<std::vector<T>>(detail::computeTotalSize(shape), value))
            , shape_(std::move(shape))
            , strides_(detail::computeRowMajorStrides(shape_))
            , offset_(0) {
        }

        // shape + initial value (initializer_list overload)
        Tensor(std::initializer_list<size_t> shape, const T& value)
            : Tensor(std::vector<size_t>(shape), value) {
        }

        // shape + flat data
        Tensor(std::vector<size_t> shape, std::vector<T> data)
            : shape_(std::move(shape))
            , strides_(detail::computeRowMajorStrides(shape_))
            , offset_(0) {
            size_t expected = detail::computeTotalSize(shape_);
            if (data.size() != expected) {
                throw DimensionError("Tensor: data size (" + std::to_string(data.size()) +
                    ") does not match shape (expected " + std::to_string(expected) + ")");
            }
            data_ = std::make_shared<std::vector<T>>(std::move(data));
        }

        // shape + flat data (initializer_list overload)
        Tensor(std::initializer_list<size_t> shape, std::vector<T> data)
            : Tensor(std::vector<size_t>(shape), std::move(data)) {
        }

        // Scalar (rank-0)
        explicit Tensor(const T& scalar)
            : data_(std::make_shared<std::vector<T>>(1, scalar))
            , shape_()
            , strides_()
            , offset_(0) {
        }

        // Conversion from Matrix<T> (rank-2)
        explicit Tensor(const Matrix<T>& matrix)
            : shape_({ matrix.rows(), matrix.cols() })
            , strides_(detail::computeRowMajorStrides(shape_))
            , offset_(0) {
            size_t total = matrix.rows() * matrix.cols();
            data_ = std::make_shared<std::vector<T>>(total);
            for (size_t i = 0; i < matrix.rows(); ++i) {
                for (size_t j = 0; j < matrix.cols(); ++j) {
                    (*data_)[i * matrix.cols() + j] = matrix(i, j);
                }
            }
        }

        // Conversion from Vector<T> (rank-1)
        explicit Tensor(const Vector<T>& vec)
            : shape_({ vec.size() })
            , strides_({ 1 })
            , offset_(0) {
            data_ = std::make_shared<std::vector<T>>(vec.size());
            for (size_t i = 0; i < vec.size(); ++i) {
                (*data_)[i] = vec[i];
            }
        }

        // Copy/move
        Tensor(const Tensor& other) = default;
        Tensor(Tensor&& other) noexcept = default;
        Tensor& operator=(const Tensor& other) = default;
        Tensor& operator=(Tensor&& other) noexcept = default;
        ~Tensor() = default;

        // ============================================================
        // Element access
        // ============================================================

        // variadic template operator() — runtime rank check
        template<typename... Indices>
            requires (std::convertible_to<Indices, size_t> && ...)
        reference operator()(Indices... indices) {
            constexpr size_t N = sizeof...(Indices);
            if (N != rank()) {
                throw DimensionError("Tensor::operator(): expected " + std::to_string(rank()) +
                    " indices, got " + std::to_string(N));
            }
            size_t idx_array[] = { static_cast<size_t>(indices)... };
            return (*data_)[computeFlatIndex(std::span<const size_t>(idx_array, N))];
        }

        template<typename... Indices>
            requires (std::convertible_to<Indices, size_t> && ...)
        const_reference operator()(Indices... indices) const {
            constexpr size_t N = sizeof...(Indices);
            if (N != rank()) {
                throw DimensionError("Tensor::operator(): expected " + std::to_string(rank()) +
                    " indices, got " + std::to_string(N));
            }
            size_t idx_array[] = { static_cast<size_t>(indices)... };
            return (*data_)[computeFlatIndex(std::span<const size_t>(idx_array, N))];
        }

        // Runtime indexed access (with bounds checking)
        reference at(std::span<const size_t> indices) {
            checkIndices(indices);
            return (*data_)[computeFlatIndex(indices)];
        }

        const_reference at(std::span<const size_t> indices) const {
            checkIndices(indices);
            return (*data_)[computeFlatIndex(indices)];
        }

        // initializer_list overload
        reference at(std::initializer_list<size_t> indices) {
            return at(std::span<const size_t>(indices.begin(), indices.size()));
        }

        const_reference at(std::initializer_list<size_t> indices) const {
            return at(std::span<const size_t>(indices.begin(), indices.size()));
        }

        // Flat index access (logical order based on strides)
        reference flat(size_t index) {
            if (index >= size()) {
                throw IndexError("Tensor::flat: index " + std::to_string(index) +
                    " out of range (size = " + std::to_string(size()) + ")");
            }
            if (isContiguous()) {
                return (*data_)[offset_ + index];
            }
            // Non-contiguous case: convert the logical index into multi-dimensional indices
            auto indices = unflattenIndex(index);
            return (*data_)[computeFlatIndex(std::span<const size_t>(indices))];
        }

        const_reference flat(size_t index) const {
            if (index >= size()) {
                throw IndexError("Tensor::flat: index " + std::to_string(index) +
                    " out of range (size = " + std::to_string(size()) + ")");
            }
            if (isContiguous()) {
                return (*data_)[offset_ + index];
            }
            auto indices = unflattenIndex(index);
            return (*data_)[computeFlatIndex(std::span<const size_t>(indices))];
        }

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

        // Rank (number of dimensions)
        size_t rank() const noexcept { return shape_.size(); }

        // Shape
        const std::vector<size_t>& shape() const noexcept { return shape_; }

        // Size along a specific axis
        size_t shape(size_t axis) const {
            if (axis >= rank()) {
                throw IndexError("Tensor::shape: axis " + std::to_string(axis) +
                    " out of range (rank = " + std::to_string(rank()) + ")");
            }
            return shape_[axis];
        }

        // Strides
        const std::vector<size_t>& strides() const noexcept { return strides_; }

        // Total element count
        size_t size() const noexcept {
            return detail::computeTotalSize(shape_);
        }

        // Whether empty (element count is 0)
        bool empty() const noexcept {
            for (auto d : shape_) {
                if (d == 0) return true;
            }
            return false;
        }

        // Whether the memory is contiguous
        bool isContiguous() const noexcept {
            return detail::isContiguousStrides(shape_, strides_);
        }

        // Whether this is a rank-0 scalar
        bool isScalar() const noexcept { return shape_.empty(); }

        // ============================================================
        // Data access
        // ============================================================

        pointer data() noexcept { return data_->data() + offset_; }
        const_pointer data() const noexcept { return data_->data() + offset_; }

        // ============================================================
        // Modification
        // ============================================================

        // Fill every element with the specified value
        void fill(const T& value) {
            ensureOwnership();
            for (size_t i = 0; i < size(); ++i) {
                flat(i) = value;
            }
        }

        // Zero out every element
        void zero() {
            fill(T(0));
        }

        // ============================================================
        // Shape changes
        // ============================================================

        // reshape: change the shape if the element count matches
        // Contiguous: view (data sharing); otherwise copy
        [[nodiscard]] Tensor reshape(std::vector<size_t> newShape) const {
            size_t newSize = detail::computeTotalSize(newShape);
            if (newSize != size()) {
                throw DimensionError("Tensor::reshape: total size mismatch (current " +
                    std::to_string(size()) + ", requested " + std::to_string(newSize) + ")");
            }

            if (isContiguous()) {
                // View: shared data
                Tensor result;
                result.data_ = data_;
                result.shape_ = std::move(newShape);
                result.strides_ = detail::computeRowMajorStrides(result.shape_);
                result.offset_ = offset_;
                return result;
            } else {
                // Non-contiguous: copy into the new shape
                Tensor result(newShape);
                for (size_t i = 0; i < size(); ++i) {
                    result.flat(i) = flat(i);
                }
                return result;
            }
        }

        // transpose: axis permutation (view: no data copy)
        [[nodiscard]] Tensor transpose(std::vector<size_t> perm) const {
            if (perm.size() != rank()) {
                throw DimensionError("Tensor::transpose: permutation size (" +
                    std::to_string(perm.size()) + ") does not match rank (" +
                    std::to_string(rank()) + ")");
            }

            // Validate that perm is a permutation of 0..rank-1
            std::vector<bool> seen(rank(), false);
            for (size_t p : perm) {
                if (p >= rank()) {
                    throw IndexError("Tensor::transpose: axis " + std::to_string(p) +
                        " out of range (rank = " + std::to_string(rank()) + ")");
                }
                if (seen[p]) {
                    throw DimensionError("Tensor::transpose: duplicate axis " + std::to_string(p));
                }
                seen[p] = true;
            }

            Tensor result;
            result.data_ = data_;
            result.offset_ = offset_;
            result.shape_.resize(rank());
            result.strides_.resize(rank());
            for (size_t i = 0; i < rank(); ++i) {
                result.shape_[i] = shape_[perm[i]];
                result.strides_[i] = strides_[perm[i]];
            }
            return result;
        }

        // Reverse-axis transpose (equivalent to ordinary matrix transpose for rank-2)
        [[nodiscard]] Tensor transpose() const {
            if (rank() == 0) return *this;
            std::vector<size_t> perm(rank());
            for (size_t i = 0; i < rank(); ++i) {
                perm[i] = rank() - 1 - i;
            }
            return transpose(perm);
        }

        // Return a contiguous copy
        [[nodiscard]] Tensor contiguous() const {
            if (isContiguous() && offset_ == 0 && data_.use_count() == 1) {
                return *this; // Already contiguous and uniquely owned
            }
            Tensor result(shape_);
            for (size_t i = 0; i < size(); ++i) {
                result.flat(i) = flat(i);
            }
            return result;
        }

        // ============================================================
        // Arithmetic operators (compound assignment)
        // ============================================================

        Tensor& operator+=(const Tensor& rhs) {
            checkSameShape(rhs, "addition");
            ensureOwnership();
            for (size_t i = 0; i < size(); ++i) {
                flat(i) = flat(i) + rhs.flat(i);
            }
            return *this;
        }

        Tensor& operator-=(const Tensor& rhs) {
            checkSameShape(rhs, "subtraction");
            ensureOwnership();
            for (size_t i = 0; i < size(); ++i) {
                flat(i) = flat(i) - rhs.flat(i);
            }
            return *this;
        }

        Tensor& operator*=(const T& scalar) {
            ensureOwnership();
            for (size_t i = 0; i < size(); ++i) {
                flat(i) = flat(i) * scalar;
            }
            return *this;
        }

        Tensor& operator/=(const T& scalar) {
            if (scalar == T(0)) {
                throw std::invalid_argument("Tensor: division by zero");
            }
            ensureOwnership();
            for (size_t i = 0; i < size(); ++i) {
                flat(i) = flat(i) / scalar;
            }
            return *this;
        }

        // ============================================================
        // Matrix/Vector conversion
        // ============================================================

        // Convert a rank-2 tensor to Matrix<T>
        [[nodiscard]] Matrix<T> toMatrix() const {
            if (rank() != 2) {
                throw DimensionError("Tensor::toMatrix: rank must be 2, got " + std::to_string(rank()));
            }
            Matrix<T> result(shape_[0], shape_[1]);
            for (size_t i = 0; i < shape_[0]; ++i) {
                for (size_t j = 0; j < shape_[1]; ++j) {
                    size_t idx[] = { i, j };
                    result(i, j) = (*data_)[computeFlatIndex(std::span<const size_t>(idx, 2))];
                }
            }
            return result;
        }

        // Convert a rank-1 tensor to Vector<T>
        [[nodiscard]] Vector<T> toVector() const {
            if (rank() != 1) {
                throw DimensionError("Tensor::toVector: rank must be 1, got " + std::to_string(rank()));
            }
            Vector<T> result(shape_[0]);
            for (size_t i = 0; i < shape_[0]; ++i) {
                size_t idx[] = { i };
                result[i] = (*data_)[computeFlatIndex(std::span<const size_t>(idx, 1))];
            }
            return result;
        }

        // Convert a rank-0 tensor to a scalar
        [[nodiscard]] T toScalar() const {
            if (rank() != 0) {
                throw DimensionError("Tensor::toScalar: rank must be 0, got " + std::to_string(rank()));
            }
            return (*data_)[offset_];
        }

        // ============================================================
        // Output
        // ============================================================

        [[nodiscard]] std::string to_string() const {
            std::ostringstream oss;
            oss << "Tensor(shape=[";
            for (size_t i = 0; i < shape_.size(); ++i) {
                if (i > 0) oss << ", ";
                oss << shape_[i];
            }
            oss << "]";

            if (isScalar()) {
                oss << ", value=" << (*data_)[offset_];
            } else if (size() <= 20) {
                // When few elements, display all of them
                oss << ", data=[";
                for (size_t i = 0; i < size(); ++i) {
                    if (i > 0) oss << ", ";
                    oss << flat(i);
                }
                oss << "]";
            } else {
                // When many elements, show only the head and tail
                oss << ", data=[";
                for (size_t i = 0; i < 5; ++i) {
                    if (i > 0) oss << ", ";
                    oss << flat(i);
                }
                oss << ", ..., ";
                for (size_t i = size() - 3; i < size(); ++i) {
                    if (i > size() - 3) oss << ", ";
                    oss << flat(i);
                }
                oss << "]";
            }
            oss << ")";
            return oss.str();
        }

        // ============================================================
        // View information (for debugging)
        // ============================================================

        // Whether the data is shared
        bool isView() const noexcept {
            return data_.use_count() > 1;
        }

        // Reference count of the shared_ptr
        long useCount() const noexcept {
            return data_.use_count();
        }

    private:
        std::shared_ptr<std::vector<T>> data_;
        std::vector<size_t> shape_;
        std::vector<size_t> strides_;
        size_t offset_ = 0;

        // ============================================================
        // Internal utilities
        // ============================================================

        // Flat-index computation from strides (no bounds checking)
        size_t computeFlatIndex(std::span<const size_t> indices) const {
            size_t flat = offset_;
            for (size_t i = 0; i < indices.size(); ++i) {
                flat += indices[i] * strides_[i];
            }
            return flat;
        }

        // Index bounds check
        void checkIndices(std::span<const size_t> indices) const {
            if (indices.size() != rank()) {
                throw DimensionError("Tensor: expected " + std::to_string(rank()) +
                    " indices, got " + std::to_string(indices.size()));
            }
            for (size_t i = 0; i < indices.size(); ++i) {
                if (indices[i] >= shape_[i]) {
                    throw IndexError("Tensor: index[" + std::to_string(i) + "] = " +
                        std::to_string(indices[i]) + " out of range (size = " +
                        std::to_string(shape_[i]) + ")");
                }
            }
        }

        // Shape equality check
        void checkSameShape(const Tensor& other, const std::string& op) const {
            if (shape_ != other.shape_) {
                throw DimensionError("Tensor " + op + ": shape mismatch");
            }
        }

        // Convert a flat index to multi-dimensional indices (C-order)
        std::vector<size_t> unflattenIndex(size_t flatIdx) const {
            std::vector<size_t> indices(rank());
            auto rowMajor = detail::computeRowMajorStrides(shape_);
            for (size_t i = 0; i < rank(); ++i) {
                indices[i] = flatIdx / rowMajor[i];
                flatIdx %= rowMajor[i];
            }
            return indices;
        }

        // For a view, materialize an independent contiguous copy
        void ensureOwnership() {
            if (data_.use_count() > 1 || !isContiguous() || offset_ != 0) {
                auto newData = std::make_shared<std::vector<T>>(size());
                for (size_t i = 0; i < size(); ++i) {
                    (*newData)[i] = flat(i);
                }
                data_ = std::move(newData);
                strides_ = detail::computeRowMajorStrides(shape_);
                offset_ = 0;
            }
        }
    };

    // ================================================================
    // Free functions: arithmetic operators
    // ================================================================

    // Tensor addition
    template<TensorScalar T>
    Tensor<T> operator+(const Tensor<T>& lhs, const Tensor<T>& rhs) {
        Tensor<T> result = lhs.contiguous();
        result += rhs;
        return result;
    }

    // Tensor subtraction
    template<TensorScalar T>
    Tensor<T> operator-(const Tensor<T>& lhs, const Tensor<T>& rhs) {
        Tensor<T> result = lhs.contiguous();
        result -= rhs;
        return result;
    }

    // Scalar multiplication (right)
    template<TensorScalar T>
    Tensor<T> operator*(const Tensor<T>& t, const T& scalar) {
        Tensor<T> result = t.contiguous();
        result *= scalar;
        return result;
    }

    // Scalar multiplication (left)
    template<TensorScalar T>
    Tensor<T> operator*(const T& scalar, const Tensor<T>& t) {
        return t * scalar;
    }

    // Scalar division
    template<TensorScalar T>
    Tensor<T> operator/(const Tensor<T>& t, const T& scalar) {
        Tensor<T> result = t.contiguous();
        result /= scalar;
        return result;
    }

    // Unary minus
    template<TensorScalar T>
    Tensor<T> operator-(const Tensor<T>& t) {
        Tensor<T> result(t.shape());
        for (size_t i = 0; i < t.size(); ++i) {
            result.flat(i) = -t.flat(i);
        }
        return result;
    }

    // ================================================================
    // Free functions: tensor operations
    // ================================================================

    // Contraction: contract along the specified axis pairs
    // Each axes[i] = {axis_of_a, axis_of_b} pair is contracted
    template<TensorScalar T>
    Tensor<T> contract(const Tensor<T>& a, const Tensor<T>& b,
        const std::vector<std::pair<size_t, size_t>>& axes) {

        // Validate axes
        for (const auto& [ax_a, ax_b] : axes) {
            if (ax_a >= a.rank()) {
                throw IndexError("contract: axis_a " + std::to_string(ax_a) +
                    " out of range (rank = " + std::to_string(a.rank()) + ")");
            }
            if (ax_b >= b.rank()) {
                throw IndexError("contract: axis_b " + std::to_string(ax_b) +
                    " out of range (rank = " + std::to_string(b.rank()) + ")");
            }
            if (a.shape(ax_a) != b.shape(ax_b)) {
                throw DimensionError("contract: dimension mismatch on axes (" +
                    std::to_string(ax_a) + ", " + std::to_string(ax_b) + "): " +
                    std::to_string(a.shape(ax_a)) + " vs " + std::to_string(b.shape(ax_b)));
            }
        }

        // Classify axes as contracted vs free
        std::vector<bool> a_contracted(a.rank(), false);
        std::vector<bool> b_contracted(b.rank(), false);
        for (const auto& [ax_a, ax_b] : axes) {
            a_contracted[ax_a] = true;
            b_contracted[ax_b] = true;
        }

        // Result shape: a's free axes + b's free axes
        std::vector<size_t> result_shape;
        std::vector<size_t> a_free_axes, b_free_axes;
        for (size_t i = 0; i < a.rank(); ++i) {
            if (!a_contracted[i]) {
                a_free_axes.push_back(i);
                result_shape.push_back(a.shape(i));
            }
        }
        for (size_t i = 0; i < b.rank(); ++i) {
            if (!b_contracted[i]) {
                b_free_axes.push_back(i);
                result_shape.push_back(b.shape(i));
            }
        }

        // Sizes of the contracted axes (product over all pairs)
        std::vector<size_t> contract_axes_a, contract_axes_b;
        std::vector<size_t> contract_sizes;
        for (const auto& [ax_a, ax_b] : axes) {
            contract_axes_a.push_back(ax_a);
            contract_axes_b.push_back(ax_b);
            contract_sizes.push_back(a.shape(ax_a));
        }

        // Case where the result becomes rank-0
        if (result_shape.empty()) {
            // All axes are contracted → scalar result
            T sum = T(0);
            // Iterate over the contracted axes
            size_t contract_total = 1;
            for (size_t s : contract_sizes) contract_total *= s;

            for (size_t ci = 0; ci < contract_total; ++ci) {
                // Unpack the contracted indices
                std::vector<size_t> cidx(contract_sizes.size());
                size_t tmp = ci;
                for (size_t k = contract_sizes.size(); k > 0; --k) {
                    cidx[k - 1] = tmp % contract_sizes[k - 1];
                    tmp /= contract_sizes[k - 1];
                }

                // Build the index for a
                std::vector<size_t> a_idx(a.rank());
                for (size_t k = 0; k < contract_axes_a.size(); ++k) {
                    a_idx[contract_axes_a[k]] = cidx[k];
                }
                // a's free axes are all empty (rank-0, so no free axes)

                // Build the index for b
                std::vector<size_t> b_idx(b.rank());
                for (size_t k = 0; k < contract_axes_b.size(); ++k) {
                    b_idx[contract_axes_b[k]] = cidx[k];
                }

                sum = sum + a.at(std::span<const size_t>(a_idx)) * b.at(std::span<const size_t>(b_idx));
            }
            return Tensor<T>(sum);
        }

        Tensor<T> result(result_shape);

        // Loop over every element of the result
        size_t result_total = result.size();
        auto result_strides = detail::computeRowMajorStrides(result_shape);

        // Total number of contracted-axis combinations
        size_t contract_total = 1;
        for (size_t s : contract_sizes) contract_total *= s;

        for (size_t ri = 0; ri < result_total; ++ri) {
            // Unpack the result's flat index into multi-dimensional indices
            std::vector<size_t> r_idx(result_shape.size());
            size_t tmp = ri;
            for (size_t k = result_shape.size(); k > 0; --k) {
                r_idx[k - 1] = tmp % result_shape[k - 1];
                tmp /= result_shape[k - 1];
            }

            // Split r_idx into a_free and b_free parts
            T sum = T(0);

            for (size_t ci = 0; ci < contract_total; ++ci) {
                // Unpack the contracted indices
                std::vector<size_t> cidx(contract_sizes.size());
                tmp = ci;
                for (size_t k = contract_sizes.size(); k > 0; --k) {
                    cidx[k - 1] = tmp % contract_sizes[k - 1];
                    tmp /= contract_sizes[k - 1];
                }

                // Build the index for a
                std::vector<size_t> a_idx(a.rank());
                for (size_t k = 0; k < a_free_axes.size(); ++k) {
                    a_idx[a_free_axes[k]] = r_idx[k];
                }
                for (size_t k = 0; k < contract_axes_a.size(); ++k) {
                    a_idx[contract_axes_a[k]] = cidx[k];
                }

                // Build the index for b
                std::vector<size_t> b_idx(b.rank());
                for (size_t k = 0; k < b_free_axes.size(); ++k) {
                    b_idx[b_free_axes[k]] = r_idx[a_free_axes.size() + k];
                }
                for (size_t k = 0; k < contract_axes_b.size(); ++k) {
                    b_idx[contract_axes_b[k]] = cidx[k];
                }

                sum = sum + a.at(std::span<const size_t>(a_idx)) * b.at(std::span<const size_t>(b_idx));
            }

            result.at(std::span<const size_t>(r_idx)) = sum;
        }

        return result;
    }

    // Tensor product (outer product): rank(result) = rank(a) + rank(b)
    template<TensorScalar T>
    Tensor<T> outerProduct(const Tensor<T>& a, const Tensor<T>& b) {
        // Result shape = a.shape + b.shape
        std::vector<size_t> result_shape;
        result_shape.reserve(a.rank() + b.rank());
        for (size_t i = 0; i < a.rank(); ++i) result_shape.push_back(a.shape(i));
        for (size_t i = 0; i < b.rank(); ++i) result_shape.push_back(b.shape(i));

        // Both rank-0: scalar product
        if (result_shape.empty()) {
            return Tensor<T>(a.toScalar() * b.toScalar());
        }

        Tensor<T> result(result_shape);

        size_t a_size = a.size();
        size_t b_size = b.size();
        for (size_t i = 0; i < a_size; ++i) {
            for (size_t j = 0; j < b_size; ++j) {
                result.flat(i * b_size + j) = a.flat(i) * b.flat(j);
            }
        }

        return result;
    }

    // Hadamard product (element-wise product)
    template<TensorScalar T>
    Tensor<T> hadamard(const Tensor<T>& a, const Tensor<T>& b) {
        if (a.shape() != b.shape()) {
            throw DimensionError("hadamard: shape mismatch");
        }
        Tensor<T> result(a.shape());
        for (size_t i = 0; i < a.size(); ++i) {
            result.flat(i) = a.flat(i) * b.flat(i);
        }
        return result;
    }

    // ================================================================
    // Factory functions
    // ================================================================

    namespace tensor {
        // Zero tensor
        template<TensorScalar T>
        Tensor<T> zeros(std::vector<size_t> shape) {
            return Tensor<T>(std::move(shape), T(0));
        }

        // Tensor of all ones
        template<TensorScalar T>
        Tensor<T> ones(std::vector<size_t> shape) {
            return Tensor<T>(std::move(shape), T(1));
        }
    }

    // ================================================================
    // Stream output
    // ================================================================

    template<TensorScalar T>
    std::ostream& operator<<(std::ostream& os, const Tensor<T>& t) {
        os << t.to_string();
        return os;
    }

    // ================================================================
    // Type traits
    // ================================================================

    // Traits that determine whether a type is a Tensor
    template<typename T>
    struct is_tensor : std::false_type {};

    template<TensorScalar T>
    struct is_tensor<Tensor<T>> : std::true_type {};

    template<typename T>
    inline constexpr bool is_tensor_v = is_tensor<T>::value;

} // namespace sangi

#endif // SANGI_TENSOR_HPP