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

// expr_templates.hpp
// Expression Templates — ベクトル・行列の遅延評価
//
// v1 + v2 + v3 のような連鎖演算で一時オブジェクト生成を排除し、
// 一回のループで全要素を計算する。
//
// 使い方:
//   Vector<double> v = v1 + 2.0 * v2 - v3 / 3.0;  // 一時 Vector 生成なし
//   Matrix<double> m = m1 + m2 - 3.0 * m3;         // 一時 Matrix 生成なし
//
// 注意: auto で式を保持するとダングリング参照の危険がある
//   auto expr = v1 + v2;   // expr は v1, v2 への参照を保持 — 寿命に注意
//   Vector<double> v = v1 + v2;  // こちらは安全 (即座に評価)

#ifndef CALX_EXPR_TEMPLATES_HPP
#define CALX_EXPR_TEMPLATES_HPP

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

#include <cstddef>
#include <concepts>
#include <functional>
#include <stdexcept>
#include <type_traits>

namespace calx {

// 式テンプレート型の基底マーカー (スカラー演算子の制約に使用)
struct ExprBase {};

// ================================================================
// ベクトル式テンプレート
// ================================================================

template<typename Derived>
struct VecExpr : ExprBase {
    const Derived& derived() const noexcept {
        return static_cast<const Derived&>(*this);
    }
};

// 単項マイナス
template<typename E>
struct VecNeg : VecExpr<VecNeg<E>> {
    using value_type = typename E::value_type;
    const E& expr_;
    explicit VecNeg(const E& e) : expr_(e) {}
    value_type operator[](std::size_t i) const { return -expr_[i]; }
    std::size_t size() const { return expr_.size(); }

    auto packet(std::size_t i) const {
        return PacketTraits<value_type>::negate(expr_.packet(i));
    }
};

// スカラー乗算
template<typename E, typename S>
struct VecScale : VecExpr<VecScale<E, S>> {
    using value_type = typename E::value_type;
    const E& expr_;
    S scalar_;
    VecScale(const E& e, S s) : expr_(e), scalar_(s) {}
    value_type operator[](std::size_t i) const {
        return static_cast<value_type>(expr_[i] * scalar_);
    }
    std::size_t size() const { return expr_.size(); }

    auto packet(std::size_t i) const {
        using PT = PacketTraits<value_type>;
        return PT::mul(expr_.packet(i), PT::set1(static_cast<value_type>(scalar_)));
    }
};

// スカラー除算 (逆数乗算に変換: div は ~20 サイクル、mul は ~4 サイクル)
template<typename E, typename S>
struct VecScaleDiv : VecExpr<VecScaleDiv<E, S>> {
    using value_type = typename E::value_type;
    const E& expr_;
    value_type inv_;  // 1/scalar を事前計算
    VecScaleDiv(const E& e, S s) : expr_(e), inv_(value_type(1) / static_cast<value_type>(s)) {}
    value_type operator[](std::size_t i) const {
        return static_cast<value_type>(expr_[i] * inv_);
    }
    std::size_t size() const { return expr_.size(); }

    auto packet(std::size_t i) const {
        using PT = PacketTraits<value_type>;
        return PT::mul(expr_.packet(i), PT::set1(inv_));
    }
};

// VecScale 型判定ヘルパー (FMA 融合用)
template<typename T> struct is_vec_scale : std::false_type {};
template<typename E, typename S> struct is_vec_scale<VecScale<E, S>> : std::true_type {};

// 二項演算ノード (加算・減算)
// Scale+Add/Sub を検出して FMA に融合: mul+add (8c) → fmadd (4c)
template<typename L, typename R, typename Op>
struct VecBinOp : VecExpr<VecBinOp<L, R, Op>> {
    using value_type = typename L::value_type;
    const L& lhs_;
    const R& rhs_;
    VecBinOp(const L& l, const R& r) : lhs_(l), rhs_(r) {
        if (l.size() != r.size())
            throw DimensionError("Vector expression: size mismatch");
    }
    value_type operator[](std::size_t i) const { return Op{}(lhs_[i], rhs_[i]); }
    std::size_t size() const { return lhs_.size(); }

    // SIMD パケット評価 — Scale+Add/Sub を FMA に融合
    auto packet(std::size_t i) const {
        using PT = PacketTraits<value_type>;
        if constexpr (std::is_same_v<Op, std::plus<>> && is_vec_scale<L>::value) {
            // s*a + b → fmadd(a, s, b)
            return PT::fmadd(lhs_.expr_.packet(i),
                             PT::set1(static_cast<value_type>(lhs_.scalar_)),
                             rhs_.packet(i));
        } else if constexpr (std::is_same_v<Op, std::plus<>> && is_vec_scale<R>::value) {
            // a + s*b → fmadd(b, s, a)
            return PT::fmadd(rhs_.expr_.packet(i),
                             PT::set1(static_cast<value_type>(rhs_.scalar_)),
                             lhs_.packet(i));
        } else if constexpr (std::is_same_v<Op, std::minus<>> && is_vec_scale<R>::value) {
            // a - s*b → fmadd(b, -s, a)
            return PT::fmadd(rhs_.expr_.packet(i),
                             PT::negate(PT::set1(static_cast<value_type>(rhs_.scalar_))),
                             lhs_.packet(i));
        } else {
            return PacketOp<Op, value_type>::apply(lhs_.packet(i), rhs_.packet(i));
        }
    }
};

// ---- ベクトル演算子 ----

template<typename L, typename R>
VecBinOp<L, R, std::plus<>> operator+(const VecExpr<L>& a, const VecExpr<R>& b) {
    return {a.derived(), b.derived()};
}

template<typename L, typename R>
VecBinOp<L, R, std::minus<>> operator-(const VecExpr<L>& a, const VecExpr<R>& b) {
    return {a.derived(), b.derived()};
}

template<typename E>
VecNeg<E> operator-(const VecExpr<E>& e) {
    return VecNeg<E>{e.derived()};
}

template<typename E, typename S>
    requires (!std::derived_from<std::remove_cvref_t<S>, ExprBase>)
VecScale<E, S> operator*(const VecExpr<E>& e, const S& s) {
    return {e.derived(), s};
}

template<typename S, typename E>
    requires (!std::derived_from<std::remove_cvref_t<S>, ExprBase>)
VecScale<E, S> operator*(const S& s, const VecExpr<E>& e) {
    return {e.derived(), s};
}

template<typename E, typename S>
    requires (!std::derived_from<std::remove_cvref_t<S>, ExprBase>)
VecScaleDiv<E, S> operator/(const VecExpr<E>& e, const S& s) {
    if (s == S{0})
        throw std::invalid_argument("Division by zero");
    return {e.derived(), s};
}

// ================================================================
// 行列式テンプレート (要素単位演算のみ)
// 行列乗算 (matrix * matrix, matrix * vector) は遅延評価しない
// ================================================================

template<typename Derived>
struct MatExpr : ExprBase {
    const Derived& derived() const noexcept {
        return static_cast<const Derived&>(*this);
    }
};

// 二項演算ノード (加算・減算)
template<typename L, typename R, typename Op>
struct MatBinOp : MatExpr<MatBinOp<L, R, Op>> {
    using value_type = typename L::value_type;
    const L& lhs_;
    const R& rhs_;
    MatBinOp(const L& l, const R& r) : lhs_(l), rhs_(r) {
        if (l.rows() != r.rows() || l.cols() != r.cols())
            throw DimensionError("Matrix expression: dimension mismatch");
    }
    value_type operator()(std::size_t i, std::size_t j) const {
        return Op{}(lhs_(i, j), rhs_(i, j));
    }
    std::size_t rows() const { return lhs_.rows(); }
    std::size_t cols() const { return lhs_.cols(); }

    // SIMD パケット評価 (線形インデックス)
    auto packet(std::size_t idx) const {
        return PacketOp<Op, value_type>::apply(lhs_.packet(idx), rhs_.packet(idx));
    }
};

// 単項マイナス
template<typename E>
struct MatNeg : MatExpr<MatNeg<E>> {
    using value_type = typename E::value_type;
    const E& expr_;
    explicit MatNeg(const E& e) : expr_(e) {}
    value_type operator()(std::size_t i, std::size_t j) const {
        return -expr_(i, j);
    }
    std::size_t rows() const { return expr_.rows(); }
    std::size_t cols() const { return expr_.cols(); }

    auto packet(std::size_t idx) const {
        return PacketTraits<value_type>::negate(expr_.packet(idx));
    }
};

// スカラー乗算
template<typename E, typename S>
struct MatScale : MatExpr<MatScale<E, S>> {
    using value_type = typename E::value_type;
    const E& expr_;
    S scalar_;
    MatScale(const E& e, S s) : expr_(e), scalar_(s) {}
    value_type operator()(std::size_t i, std::size_t j) const {
        return static_cast<value_type>(expr_(i, j) * scalar_);
    }
    std::size_t rows() const { return expr_.rows(); }
    std::size_t cols() const { return expr_.cols(); }

    auto packet(std::size_t idx) const {
        using PT = PacketTraits<value_type>;
        return PT::mul(expr_.packet(idx), PT::set1(static_cast<value_type>(scalar_)));
    }
};

// スカラー除算 (逆数乗算に変換)
template<typename E, typename S>
struct MatScaleDiv : MatExpr<MatScaleDiv<E, S>> {
    using value_type = typename E::value_type;
    const E& expr_;
    value_type inv_;
    MatScaleDiv(const E& e, S s) : expr_(e), inv_(value_type(1) / static_cast<value_type>(s)) {}
    value_type operator()(std::size_t i, std::size_t j) const {
        return static_cast<value_type>(expr_(i, j) * inv_);
    }
    std::size_t rows() const { return expr_.rows(); }
    std::size_t cols() const { return expr_.cols(); }

    auto packet(std::size_t idx) const {
        using PT = PacketTraits<value_type>;
        return PT::mul(expr_.packet(idx), PT::set1(inv_));
    }
};

// ---- 行列演算子 (要素単位) ----

template<typename L, typename R>
MatBinOp<L, R, std::plus<>> operator+(const MatExpr<L>& a, const MatExpr<R>& b) {
    return {a.derived(), b.derived()};
}

template<typename L, typename R>
MatBinOp<L, R, std::minus<>> operator-(const MatExpr<L>& a, const MatExpr<R>& b) {
    return {a.derived(), b.derived()};
}

template<typename E>
MatNeg<E> operator-(const MatExpr<E>& e) {
    return MatNeg<E>{e.derived()};
}

template<typename E, typename S>
    requires (!std::derived_from<std::remove_cvref_t<S>, ExprBase>)
MatScale<E, S> operator*(const MatExpr<E>& e, const S& s) {
    return {e.derived(), s};
}

template<typename S, typename E>
    requires (!std::derived_from<std::remove_cvref_t<S>, ExprBase>)
MatScale<E, S> operator*(const S& s, const MatExpr<E>& e) {
    return {e.derived(), s};
}

template<typename E, typename S>
    requires (!std::derived_from<std::remove_cvref_t<S>, ExprBase>)
MatScaleDiv<E, S> operator/(const MatExpr<E>& e, const S& s) {
    if (s == S{0})
        throw std::invalid_argument("Division by zero");
    return {e.derived(), s};
}

} // namespace calx

#endif // CALX_EXPR_TEMPLATES_HPP