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

// Tensor.hpp
// 任意階数テンソルクラス
//
// 特徴:
//   - 動的ランク・動的シェイプ
//   - ストライドベースのビューセマンティクス (transpose/reshape がデータコピー不要)
//   - shared_ptr によるデータ共有 (ビューが安全に元データを参照)
//   - Matrix<T>, Vector<T> との相互変換
//   - C++23 concepts による型制約

#ifndef CALX_TENSOR_HPP
#define CALX_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 calx {

    // 前方宣言
    template<typename T> class Matrix;
    template<typename T> class Vector;

    // ================================================================
    // 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) };
    };

    // ================================================================
    // ストライド計算ユーティリティ (detail)
    // ================================================================

    namespace detail {

        // C-order (row-major) ストライドの計算
        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;
        }

        // shape から全要素数を計算
        inline size_t computeTotalSize(const std::vector<size_t>& shape) {
            if (shape.empty()) return 1; // rank-0 スカラーは要素数 1
            return std::accumulate(shape.begin(), shape.end(), size_t(1), std::multiplies<size_t>());
        }

        // ストライドがコンティギュアス (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> クラス
    // ================================================================

    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*;

        // ============================================================
        // コンストラクタ
        // ============================================================

        // デフォルト: rank-0 スカラー (値 = T(0))
        Tensor()
            : data_(std::make_shared<std::vector<T>>(1, T(0)))
            , shape_()
            , strides_()
            , offset_(0) {
        }

        // shape 指定 (ゼロ初期化)
        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) {
        }

        // shape 指定 (initializer_list 版, オーバーロード解決用)
        explicit Tensor(std::initializer_list<size_t> shape)
            : Tensor(std::vector<size_t>(shape)) {
        }

        // shape + 初期値
        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 + 初期値 (initializer_list 版)
        Tensor(std::initializer_list<size_t> shape, const T& value)
            : Tensor(std::vector<size_t>(shape), value) {
        }

        // shape + フラットデータ
        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 + フラットデータ (initializer_list 版)
        Tensor(std::initializer_list<size_t> shape, std::vector<T> data)
            : Tensor(std::vector<size_t>(shape), std::move(data)) {
        }

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

        // 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);
                }
            }
        }

        // 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];
            }
        }

        // コピー/ムーブ
        Tensor(const Tensor& other) = default;
        Tensor(Tensor&& other) noexcept = default;
        Tensor& operator=(const Tensor& other) = default;
        Tensor& operator=(Tensor&& other) noexcept = default;
        ~Tensor() = default;

        // ============================================================
        // 要素アクセス
        // ============================================================

        // variadic template operator() — ランタイム rank チェック
        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))];
        }

        // ランタイムインデックスアクセス (境界チェック付き)
        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 版
        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()));
        }

        // フラットインデックスアクセス (ストライドに基づく論理順序)
        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];
            }
            // 非コンティギュアスの場合: 論理インデックスを多次元に変換
            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))];
        }

        // ============================================================
        // プロパティ
        // ============================================================

        // 階数 (次元数)
        size_t rank() const noexcept { return shape_.size(); }

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

        // 特定軸のサイズ
        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];
        }

        // ストライド
        const std::vector<size_t>& strides() const noexcept { return strides_; }

        // 全要素数
        size_t size() const noexcept {
            return detail::computeTotalSize(shape_);
        }

        // 空かどうか (要素数 0)
        bool empty() const noexcept {
            for (auto d : shape_) {
                if (d == 0) return true;
            }
            return false;
        }

        // メモリが連続かどうか
        bool isContiguous() const noexcept {
            return detail::isContiguousStrides(shape_, strides_);
        }

        // rank-0 スカラーかどうか
        bool isScalar() const noexcept { return shape_.empty(); }

        // ============================================================
        // データアクセス
        // ============================================================

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

        // ============================================================
        // 変更操作
        // ============================================================

        // 全要素を指定値で埋める
        void fill(const T& value) {
            ensureOwnership();
            for (size_t i = 0; i < size(); ++i) {
                flat(i) = value;
            }
        }

        // 全要素をゼロにする
        void zero() {
            fill(T(0));
        }

        // ============================================================
        // 形状変更
        // ============================================================

        // reshape: 要素数が一致すれば形状を変更
        // contiguous ならビュー (データ共有), そうでなければコピー
        [[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()) {
                // ビュー: データ共有
                Tensor result;
                result.data_ = data_;
                result.shape_ = std::move(newShape);
                result.strides_ = detail::computeRowMajorStrides(result.shape_);
                result.offset_ = offset_;
                return result;
            } else {
                // 非コンティギュアス: コピーして新しい形状に
                Tensor result(newShape);
                for (size_t i = 0; i < size(); ++i) {
                    result.flat(i) = flat(i);
                }
                return result;
            }
        }

        // transpose: 軸の置換 (ビュー: データコピーなし)
        [[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()) + ")");
            }

            // 置換の妥当性チェック (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;
        }

        // 全軸反転の転置 (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);
        }

        // コンティギュアスなコピーを返す
        [[nodiscard]] Tensor contiguous() const {
            if (isContiguous() && offset_ == 0 && data_.use_count() == 1) {
                return *this; // 既にコンティギュアスで単独所有
            }
            Tensor result(shape_);
            for (size_t i = 0; i < size(); ++i) {
                result.flat(i) = flat(i);
            }
            return result;
        }

        // ============================================================
        // 算術演算子 (複合代入)
        // ============================================================

        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 変換
        // ============================================================

        // rank-2 テンソルを 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;
        }

        // rank-1 テンソルを 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;
        }

        // rank-0 テンソルをスカラーに変換
        [[nodiscard]] T toScalar() const {
            if (rank() != 0) {
                throw DimensionError("Tensor::toScalar: rank must be 0, got " + std::to_string(rank()));
            }
            return (*data_)[offset_];
        }

        // ============================================================
        // 出力
        // ============================================================

        [[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) {
                // 要素数が少ない場合は全要素を表示
                oss << ", data=[";
                for (size_t i = 0; i < size(); ++i) {
                    if (i > 0) oss << ", ";
                    oss << flat(i);
                }
                oss << "]";
            } else {
                // 要素数が多い場合は先頭と末尾のみ
                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();
        }

        // ============================================================
        // ビュー情報 (デバッグ用)
        // ============================================================

        // データが共有されているかどうか
        bool isView() const noexcept {
            return data_.use_count() > 1;
        }

        // 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;

        // ============================================================
        // 内部ユーティリティ
        // ============================================================

        // ストライドに基づくフラットインデックス計算 (チェックなし)
        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;
        }

        // インデックス境界チェック
        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 一致チェック
        void checkSameShape(const Tensor& other, const std::string& op) const {
            if (shape_ != other.shape_) {
                throw DimensionError("Tensor " + op + ": shape mismatch");
            }
        }

        // フラットインデックスを多次元インデックスに変換 (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;
        }

        // ビューの場合、独立したコンティギュアスコピーを作成する
        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;
            }
        }
    };

    // ================================================================
    // フリー関数: 算術演算子
    // ================================================================

    // テンソル加算
    template<TensorScalar T>
    Tensor<T> operator+(const Tensor<T>& lhs, const Tensor<T>& rhs) {
        Tensor<T> result = lhs.contiguous();
        result += rhs;
        return result;
    }

    // テンソル減算
    template<TensorScalar T>
    Tensor<T> operator-(const Tensor<T>& lhs, const Tensor<T>& rhs) {
        Tensor<T> result = lhs.contiguous();
        result -= rhs;
        return result;
    }

    // スカラー乗算 (右)
    template<TensorScalar T>
    Tensor<T> operator*(const Tensor<T>& t, const T& scalar) {
        Tensor<T> result = t.contiguous();
        result *= scalar;
        return result;
    }

    // スカラー乗算 (左)
    template<TensorScalar T>
    Tensor<T> operator*(const T& scalar, const Tensor<T>& t) {
        return t * scalar;
    }

    // スカラー除算
    template<TensorScalar T>
    Tensor<T> operator/(const Tensor<T>& t, const T& scalar) {
        Tensor<T> result = t.contiguous();
        result /= scalar;
        return result;
    }

    // 単項マイナス
    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;
    }

    // ================================================================
    // フリー関数: テンソル演算
    // ================================================================

    // 縮約 (contraction): 指定した軸ペアで縮約
    // axes[i] = {axis_of_a, axis_of_b} の各ペアを縮約する
    template<TensorScalar T>
    Tensor<T> contract(const Tensor<T>& a, const Tensor<T>& b,
        const std::vector<std::pair<size_t, size_t>>& 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)));
            }
        }

        // 縮約軸と自由軸を分類
        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;
        }

        // 結果の shape: a の自由軸 + b の自由軸
        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));
            }
        }

        // 縮約軸のサイズ (全ペアの積)
        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));
        }

        // 結果が rank-0 になる場合
        if (result_shape.empty()) {
            // 全軸が縮約される → スカラー結果
            T sum = T(0);
            // 縮約軸のイテレーション
            size_t contract_total = 1;
            for (size_t s : contract_sizes) contract_total *= s;

            for (size_t ci = 0; ci < contract_total; ++ci) {
                // 縮約インデックスを展開
                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];
                }

                // 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 の自由軸は全て空 → rank-0 のため自由軸なし

                // 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);

        // 結果の全要素についてループ
        size_t result_total = result.size();
        auto result_strides = detail::computeRowMajorStrides(result_shape);

        // 縮約軸の全組み合わせ数
        size_t contract_total = 1;
        for (size_t s : contract_sizes) contract_total *= s;

        for (size_t ri = 0; ri < result_total; ++ri) {
            // 結果のフラットインデックスを多次元に展開
            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];
            }

            // r_idx を a_free, b_free に分離
            T sum = T(0);

            for (size_t ci = 0; ci < contract_total; ++ci) {
                // 縮約インデックスを展開
                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];
                }

                // 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];
                }

                // 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;
    }

    // テンソル積 (外積): rank(result) = rank(a) + rank(b)
    template<TensorScalar T>
    Tensor<T> outerProduct(const Tensor<T>& a, const Tensor<T>& b) {
        // 結果の 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));

        // rank-0 同士の場合: スカラー積
        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 積 (要素ごとの積)
    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;
    }

    // ================================================================
    // ファクトリ関数
    // ================================================================

    namespace tensor {
        // ゼロテンソル
        template<TensorScalar T>
        Tensor<T> zeros(std::vector<size_t> shape) {
            return Tensor<T>(std::move(shape), T(0));
        }

        // 全要素 1 のテンソル
        template<TensorScalar T>
        Tensor<T> ones(std::vector<size_t> shape) {
            return Tensor<T>(std::move(shape), T(1));
        }
    }

    // ================================================================
    // ストリーム出力
    // ================================================================

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

    // ================================================================
    // 型特性
    // ================================================================

    // Tensor かどうかを判定する traits
    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 calx

#endif // CALX_TENSOR_HPP
