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

// expr_templates.hpp
// Expression templates — lazy evaluation for vectors and matrices
//
// Eliminates temporary objects from chained operations like v1 + v2 + v3, computing
// every element in a single loop.
//
// Usage:
//   Vector<double> v = v1 + 2.0 * v2 - v3 / 3.0;  // no intermediate Vector created
//   Matrix<double> m = m1 + m2 - 3.0 * m3;         // no intermediate Matrix created
//
// Caution: holding an expression with auto risks dangling references.
//   auto expr = v1 + v2;          // expr holds references to v1, v2 — mind their lifetime
//   Vector<double> v = v1 + v2;   // safe (evaluated immediately)

#ifndef SANGI_EXPR_TEMPLATES_HPP
#define SANGI_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 sangi {

// Forward declarations of Matrix / Vector (so expr_operand_store can specialize within expression templates)
template<typename T> class Matrix;
template<typename T> class Vector;

// Base marker for expression-template types (used to constrain scalar operators)
struct ExprBase {};

// ================================================================
// Operand-storage selection (preserves value category)
// ================================================================
// Because Matrix * Matrix is evaluated eagerly (returning a Matrix), prvalue
// Matrices may appear as operands inside an ET expression:
//     Matrix m = A*B + C*D + E;    // A*B and C*D are prvalue Matrices
//
// The previous implementation accepted operator+(const MatExpr<L>&, const MatExpr<R>&)
// and MatBinOp stored const L&, so prvalue Matrix temporaries ended up as
// dangling references inside the ET chain.
// (Depending on heap fragmentation, correct values sometimes came out, and we
// observed nondeterministic n>=20 residuals jumping to the order of ||Q||.)
//
// Fix: use forwarding references at operator boundaries to preserve value category.
// Additionally, wrap lvalues in std::reference_wrapper<const T> for storage
// (raw const T& members triggered UB under MSVC /O2 inlining).
//   - lvalue input → std::reference_wrapper<const T> (zero copy, value type)
//   - rvalue input → T (owned by the ET node)
template<class T>
using expr_store_t =
    std::conditional_t<std::is_lvalue_reference_v<T>,
                       std::reference_wrapper<const std::remove_reference_t<T>>,
                       std::remove_cvref_t<T>>;

// Extract the stored operand as a reference (used inside operator() and packet())
// For reference_wrapper, call .get(); otherwise return the value as-is
template<class T>
const T& expr_unwrap(const std::reference_wrapper<const T>& x) noexcept { return x.get(); }

template<class T>
    requires (std::derived_from<T, ExprBase>)
const T& expr_unwrap(const T& x) noexcept { return x; }

// Factory that converts an operand into its storage type
template<class T>
auto make_expr_store(T&& x) {
    if constexpr (std::is_lvalue_reference_v<T>) {
        return std::cref(x);
    } else {
        return std::remove_cvref_t<T>(std::forward<T>(x));
    }
}

// Obtain the underlying expression type from the store type (T or std::reference_wrapper<const U>)
template<class Store>
struct expr_underlying { using type = std::remove_cvref_t<Store>; };

template<class U>
struct expr_underlying<std::reference_wrapper<const U>> { using type = U; };

template<class U>
struct expr_underlying<std::reference_wrapper<U>> { using type = U; };

template<class Store>
using expr_underlying_t = typename expr_underlying<std::remove_cvref_t<Store>>::type;

template<class T>
concept ExprLike = std::derived_from<std::remove_cvref_t<T>, ExprBase>;

// ================================================================
// Vector expression templates
// ================================================================

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

// Concept for VecExpr-derived types (checks whether VecExpr is a base class)
template<class T>
concept VecExprLike = std::derived_from<std::remove_cvref_t<T>, ExprBase>
                   && requires(std::remove_cvref_t<T> const& v) { v.size(); v[0]; };

// Detect matrix-like types (those exposing rows()/cols()) — prevents Matrix/SparseMatrix
// from being passed as VecScale's S and being interpreted as value_type * Matrix.
namespace detail {
    template<typename T, typename = void>
    struct has_matrix_shape : std::false_type {};
    template<typename T>
    struct has_matrix_shape<T, std::void_t<
        decltype(std::declval<const T&>().rows()),
        decltype(std::declval<const T&>().cols())>> : std::true_type {};
}

// Constraint on S for VecScale: not derived from ExprBase and not matrix-like (i.e., scalar-like)
template<class T>
concept ScalarLike = !std::derived_from<std::remove_cvref_t<T>, ExprBase>
                  && !detail::has_matrix_shape<std::remove_cvref_t<T>>::value;

// Unary minus
template<typename EStore>
struct VecNeg : VecExpr<VecNeg<EStore>> {
    using Expr = expr_underlying_t<EStore>;
    using value_type = typename Expr::value_type;
    EStore expr_;
    template<class E>
    explicit VecNeg(E&& e) : expr_(make_expr_store<E>(std::forward<E>(e))) {}
    value_type operator[](std::size_t i) const { return -expr_unwrap(expr_)[i]; }
    std::size_t size() const { return expr_unwrap(expr_).size(); }

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

// Scalar multiplication
template<typename EStore, typename S>
struct VecScale : VecExpr<VecScale<EStore, S>> {
    using Expr = expr_underlying_t<EStore>;
    using value_type = typename Expr::value_type;
    EStore expr_;
    S scalar_;
    template<class E>
    VecScale(E&& e, S s) : expr_(make_expr_store<E>(std::forward<E>(e))), scalar_(s) {}
    value_type operator[](std::size_t i) const {
        return static_cast<value_type>(expr_unwrap(expr_)[i] * scalar_);
    }
    std::size_t size() const { return expr_unwrap(expr_).size(); }

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

// Scalar division (converted to reciprocal multiplication; div is ~20 cycles, mul is ~4 cycles)
template<typename EStore, typename S>
struct VecScaleDiv : VecExpr<VecScaleDiv<EStore, S>> {
    using Expr = expr_underlying_t<EStore>;
    using value_type = typename Expr::value_type;
    EStore expr_;
    value_type inv_;  // 1/scalar precomputed
    template<class E>
    VecScaleDiv(E&& e, S s)
        : expr_(make_expr_store<E>(std::forward<E>(e))),
          inv_(value_type(1) / static_cast<value_type>(s)) {}
    value_type operator[](std::size_t i) const {
        return static_cast<value_type>(expr_unwrap(expr_)[i] * inv_);
    }
    std::size_t size() const { return expr_unwrap(expr_).size(); }

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

// Helper trait that detects VecScale (used for FMA fusion)
template<typename T> struct is_vec_scale : std::false_type {};
template<typename EStore, typename S> struct is_vec_scale<VecScale<EStore, S>> : std::true_type {};

// Binary-operation node (addition / subtraction)
// Detects Scale+Add/Sub and fuses them into FMA: mul+add (8c) → fmadd (4c)
template<typename LStore, typename RStore, typename Op>
struct VecBinOp : VecExpr<VecBinOp<LStore, RStore, Op>> {
    using LExpr = expr_underlying_t<LStore>;
    using RExpr = expr_underlying_t<RStore>;
    using value_type = typename LExpr::value_type;
    LStore lhs_;
    RStore rhs_;
    template<class L, class R>
    VecBinOp(L&& l, R&& r)
        : lhs_(make_expr_store<L>(std::forward<L>(l))),
          rhs_(make_expr_store<R>(std::forward<R>(r))) {
        if (expr_unwrap(lhs_).size() != expr_unwrap(rhs_).size())
            throw DimensionError("Vector expression: size mismatch");
    }
    value_type operator[](std::size_t i) const {
        return Op{}(expr_unwrap(lhs_)[i], expr_unwrap(rhs_)[i]);
    }
    std::size_t size() const { return expr_unwrap(lhs_).size(); }

    // SIMD packet evaluation — fuse Scale+Add/Sub into FMA
    auto packet(std::size_t i) const {
        using PT = PacketTraits<value_type>;
        const auto& lhs = expr_unwrap(lhs_);
        const auto& rhs = expr_unwrap(rhs_);
        if constexpr (std::is_same_v<Op, std::plus<>> && is_vec_scale<LExpr>::value) {
            // s*a + b → fmadd(a, s, b)
            return PT::fmadd(expr_unwrap(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<RExpr>::value) {
            // a + s*b → fmadd(b, s, a)
            return PT::fmadd(expr_unwrap(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<RExpr>::value) {
            // a - s*b → fmadd(b, -s, a)
            return PT::fmadd(expr_unwrap(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));
        }
    }
};

// ---- Vector operators (forwarding references preserve value category) ----

template<class L, class R>
    requires (VecExprLike<L> && VecExprLike<R>)
auto operator+(L&& a, R&& b) {
    using LS = expr_store_t<L&&>;
    using RS = expr_store_t<R&&>;
    return VecBinOp<LS, RS, std::plus<>>(std::forward<L>(a), std::forward<R>(b));
}

template<class L, class R>
    requires (VecExprLike<L> && VecExprLike<R>)
auto operator-(L&& a, R&& b) {
    using LS = expr_store_t<L&&>;
    using RS = expr_store_t<R&&>;
    return VecBinOp<LS, RS, std::minus<>>(std::forward<L>(a), std::forward<R>(b));
}

template<class E>
    requires VecExprLike<E>
auto operator-(E&& e) {
    using ES = expr_store_t<E&&>;
    return VecNeg<ES>(std::forward<E>(e));
}

/// Element-wise subtraction scalar - vector (s - v[i]).
/// No dedicated ET node exists yet, so the result is materialized into a Vector<T> and returned.
template<class S, class E>
    requires (ScalarLike<S> && VecExprLike<E>)
auto operator-(const S& s, const E& e) {
    using T = std::remove_cvref_t<decltype(e[0])>;
    Vector<T> r(e.size());
    for (std::size_t i = 0; i < e.size(); ++i) r[i] = T(s) - e[i];
    return r;
}

template<class E, class S>
    requires (VecExprLike<E> && ScalarLike<S>)
auto operator*(E&& e, const S& s) {
    using ES = expr_store_t<E&&>;
    return VecScale<ES, S>(std::forward<E>(e), s);
}

template<class S, class E>
    requires (VecExprLike<E> && ScalarLike<S>)
auto operator*(const S& s, E&& e) {
    using ES = expr_store_t<E&&>;
    return VecScale<ES, S>(std::forward<E>(e), s);
}

template<class E, class S>
    requires (VecExprLike<E> && ScalarLike<S>)
auto operator/(E&& e, const S& s) {
    if (s == S{0})
        throw std::invalid_argument("Division by zero");
    using ES = expr_store_t<E&&>;
    return VecScaleDiv<ES, S>(std::forward<E>(e), s);
}

// ================================================================
// Matrix expression templates (element-wise operations only)
// Matrix multiplication (matrix * matrix, matrix * vector) is not lazy
// ================================================================

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

// Concept for MatExpr-derived types
template<class T>
concept MatExprLike = std::derived_from<std::remove_cvref_t<T>, ExprBase>
                   && requires(std::remove_cvref_t<T> const& m) { m.rows(); m.cols(); m(0, 0); };

// Binary-operation node (addition / subtraction)
template<typename LStore, typename RStore, typename Op>
struct MatBinOp : MatExpr<MatBinOp<LStore, RStore, Op>> {
    using LExpr = expr_underlying_t<LStore>;
    using RExpr = expr_underlying_t<RStore>;
    using value_type = typename LExpr::value_type;
    LStore lhs_;
    RStore rhs_;
    template<class L, class R>
    MatBinOp(L&& l, R&& r)
        : lhs_(make_expr_store<L>(std::forward<L>(l))),
          rhs_(make_expr_store<R>(std::forward<R>(r))) {
        const auto& lhs = expr_unwrap(lhs_);
        const auto& rhs = expr_unwrap(rhs_);
        if (lhs.rows() != rhs.rows() || lhs.cols() != rhs.cols())
            throw DimensionError("Matrix expression: dimension mismatch");
    }
    value_type operator()(std::size_t i, std::size_t j) const {
        return Op{}(expr_unwrap(lhs_)(i, j), expr_unwrap(rhs_)(i, j));
    }
    std::size_t rows() const { return expr_unwrap(lhs_).rows(); }
    std::size_t cols() const { return expr_unwrap(lhs_).cols(); }

    // SIMD packet evaluation (linear index)
    auto packet(std::size_t idx) const {
        return PacketOp<Op, value_type>::apply(
            expr_unwrap(lhs_).packet(idx),
            expr_unwrap(rhs_).packet(idx));
    }
};

// Unary minus
template<typename EStore>
struct MatNeg : MatExpr<MatNeg<EStore>> {
    using Expr = expr_underlying_t<EStore>;
    using value_type = typename Expr::value_type;
    EStore expr_;
    template<class E>
    explicit MatNeg(E&& e) : expr_(make_expr_store<E>(std::forward<E>(e))) {}
    value_type operator()(std::size_t i, std::size_t j) const {
        return -expr_unwrap(expr_)(i, j);
    }
    std::size_t rows() const { return expr_unwrap(expr_).rows(); }
    std::size_t cols() const { return expr_unwrap(expr_).cols(); }

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

// Scalar multiplication
template<typename EStore, typename S>
struct MatScale : MatExpr<MatScale<EStore, S>> {
    using Expr = expr_underlying_t<EStore>;
    using value_type = typename Expr::value_type;
    EStore expr_;
    S scalar_;
    template<class E>
    MatScale(E&& e, S s) : expr_(make_expr_store<E>(std::forward<E>(e))), scalar_(s) {}
    value_type operator()(std::size_t i, std::size_t j) const {
        return static_cast<value_type>(expr_unwrap(expr_)(i, j) * scalar_);
    }
    std::size_t rows() const { return expr_unwrap(expr_).rows(); }
    std::size_t cols() const { return expr_unwrap(expr_).cols(); }

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

// Scalar division (converted to reciprocal multiplication)
template<typename EStore, typename S>
struct MatScaleDiv : MatExpr<MatScaleDiv<EStore, S>> {
    using Expr = expr_underlying_t<EStore>;
    using value_type = typename Expr::value_type;
    EStore expr_;
    value_type inv_;
    template<class E>
    MatScaleDiv(E&& e, S s)
        : expr_(make_expr_store<E>(std::forward<E>(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_unwrap(expr_)(i, j) * inv_);
    }
    std::size_t rows() const { return expr_unwrap(expr_).rows(); }
    std::size_t cols() const { return expr_unwrap(expr_).cols(); }

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

// ---- Matrix operators (element-wise, forwarding references preserve value category) ----

template<class L, class R>
    requires (MatExprLike<L> && MatExprLike<R>)
auto operator+(L&& a, R&& b) {
    using LS = expr_store_t<L&&>;
    using RS = expr_store_t<R&&>;
    return MatBinOp<LS, RS, std::plus<>>(std::forward<L>(a), std::forward<R>(b));
}

template<class L, class R>
    requires (MatExprLike<L> && MatExprLike<R>)
auto operator-(L&& a, R&& b) {
    using LS = expr_store_t<L&&>;
    using RS = expr_store_t<R&&>;
    return MatBinOp<LS, RS, std::minus<>>(std::forward<L>(a), std::forward<R>(b));
}

template<class E>
    requires MatExprLike<E>
auto operator-(E&& e) {
    using ES = expr_store_t<E&&>;
    return MatNeg<ES>(std::forward<E>(e));
}

template<class E, class S>
    requires (MatExprLike<E> && ScalarLike<S>)
auto operator*(E&& e, const S& s) {
    using ES = expr_store_t<E&&>;
    return MatScale<ES, S>(std::forward<E>(e), s);
}

template<class S, class E>
    requires (MatExprLike<E> && ScalarLike<S>)
auto operator*(const S& s, E&& e) {
    using ES = expr_store_t<E&&>;
    return MatScale<ES, S>(std::forward<E>(e), s);
}

template<class E, class S>
    requires (MatExprLike<E> && !std::derived_from<std::remove_cvref_t<S>, ExprBase>)
auto operator/(E&& e, const S& s) {
    if (s == S{0})
        throw std::invalid_argument("Division by zero");
    using ES = expr_store_t<E&&>;
    return MatScaleDiv<ES, S>(std::forward<E>(e), s);
}

} // namespace sangi

#endif // SANGI_EXPR_TEMPLATES_HPP