Arithmetic.hpp

#pragma once
 
#include <stratax/core/dtypes/Concepts.hpp>
#include <stratax/exceptions/Exceptions.hpp>
#include <stratax/ops/Broadcasting.hpp>
 
#include <functional>

template<Array L, Array R, typename Op>
requires (
    Numeric<typename L::value_type> &&
    Numeric<typename R::value_type>
)
auto binary_op(
    const L& lhs,
    const R& rhs,
    Op op,
    bool check_zero_divisor = false)
{
    auto checked_op = [&](const auto& left, const auto& right)
    {
        if (check_zero_divisor &&
            right == typename R::value_type{})
        {
            throw Exceptions::ZeroDivisionError(
                "Array division divisor element cannot be zero.");
        }
 
        return op(left, right);
    };
 
    using result_value_type =
        stratax::core::promote_t<
            typename L::value_type,
            typename R::value_type>;
 
    return stratax::core::broadcasted_op<result_value_type>(
        lhs,
        rhs,
        checked_op);
}

template<Array A, Numeric Scalar, typename Op>
auto binary_scalar_op(const A& lhs, const Scalar& rhs, Op op, bool check_zero_divisor = false)
{
    if (check_zero_divisor && rhs == Scalar{})
    {
        throw Exceptions::ZeroDivisionError("Array division scalar divisor cannot be zero.");
    }
 
    return broadcasted_op(lhs, rhs, op);
}

template<Numeric Scalar, Array A, typename Op>
auto binary_scalar_op(const Scalar& lhs, const A& rhs, Op op, bool check_zero_divisor = false)
{
    auto checked_op = [&](const auto& left, const auto& right)
    {
        if (check_zero_divisor && right == typename A::value_type{})
        {
            throw Exceptions::ZeroDivisionError("Scalar division divisor element cannot be zero.");
        }
 
        return op(left, right);
    };
 
    return broadcasted_op(lhs, rhs, checked_op);
}

template<Array L, Array R, typename Op>
requires (
    Numeric<typename L::value_type> &&
    Numeric<typename R::value_type>
)
L& compound_op(
    L& lhs,
    const R& rhs,
    Op op,
    bool check_zero_divisor = false)
{
    const auto result_shape =
        broadcasted_shape(lhs.shape(), rhs.shape());
 
    // Compound assignment cannot change the lhs shape.
    if (result_shape != lhs.shape())
    {
        throw Exceptions::BroadcastError(
            "Compound assignment cannot change the left-hand shape.");
    }
 
    for (std::size_t i = 0; i < lhs.size(); ++i)
    {
        const std::size_t rhs_index =
            stratax::core::broadcast_detail::flat_operand_index(
                i,
                lhs.shape(),
                rhs.shape());
 
        if (check_zero_divisor &&
            rhs[rhs_index] == typename R::value_type{})
        {
            throw Exceptions::ZeroDivisionError(
                "Compound division divisor element cannot be zero.");
        }
 
        lhs[i] = static_cast<typename L::value_type>(
            op(lhs[i], rhs[rhs_index]));
    }
 
    return lhs;
}

template<Array A, Numeric S, typename Op>
A& compound_scalar_op(
    A& lhs,
    const S& rhs,
    Op op,
    bool check_zero_divisor = false)
{
    if (check_zero_divisor && rhs == S{})
    {
        throw Exceptions::ZeroDivisionError(
            "Compound division scalar divisor cannot be zero.");
    }
 
    for (std::size_t i = 0; i < lhs.size(); ++i)
    {
        lhs[i] = static_cast<typename A::value_type>(
            op(lhs[i], rhs));
    }
 
    return lhs;
}

template<Array L, Array R>
requires (
    Numeric<typename L::value_type> &&
    Numeric<typename R::value_type>
)
auto operator+(const L& lhs, const R& rhs)
{
    return binary_op(lhs, rhs, std::plus<>{});
}

template<Array L, Array R>
requires (
    Numeric<typename L::value_type> &&
    Numeric<typename R::value_type>
)
auto operator-(const L& lhs, const R& rhs)
{
    return binary_op(lhs, rhs, std::minus<>{});
}

template<Array L, Array R>
requires (
    Numeric<typename L::value_type> &&
    Numeric<typename R::value_type>
)
auto operator*(const L& lhs, const R& rhs)
{
    return binary_op(lhs, rhs, std::multiplies<>{});
}

template<Array L, Array R>
requires (
    Numeric<typename L::value_type> &&
    Numeric<typename R::value_type>
)
auto operator/(const L& lhs, const R& rhs)
{
    return binary_op(lhs, rhs, std::divides<>{}, true);
}

template<Array A, Numeric Scalar>
requires Numeric<typename A::value_type>
auto operator+(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_op(lhs, rhs, std::plus<>{});
}

template<Array A, Numeric Scalar>
requires Numeric<typename A::value_type>
auto operator-(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_op(lhs, rhs, std::minus<>{});
}

template<Array A, Numeric Scalar>
requires Numeric<typename A::value_type>
auto operator*(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_op(lhs, rhs, std::multiplies<>{});
}

template<Array A, Numeric Scalar>
requires Numeric<typename A::value_type>
auto operator/(const A& lhs, const Scalar& rhs)
{
    return binary_scalar_op(lhs, rhs, std::divides<>{}, true);
}

template<Numeric Scalar, Array A>
requires Numeric<typename A::value_type>
auto operator+(const Scalar& lhs, const A& rhs)
{
    return rhs + lhs;
}

template<Numeric Scalar, Array A>
auto operator-(const Scalar& lhs, const A& rhs)
{
    return binary_scalar_op(lhs, rhs, std::minus<>{});
}

template<Numeric Scalar, Array A>
requires Numeric<typename A::value_type>
auto operator*(const Scalar& lhs, const A& rhs)
{
    return rhs * lhs;
}

template<Numeric Scalar, Array A>
requires Numeric<typename A::value_type>
auto operator/(const Scalar& lhs, const A& rhs)
{
    return binary_scalar_op(lhs, rhs, std::divides<>{}, true);
}

template<Array L, Array R>
requires (
    Numeric<typename L::value_type> &&
    Numeric<typename R::value_type>
)
L& operator+=(L& lhs, const R& rhs)
{
    return compound_op(lhs, rhs, std::plus<>{});
}

template<Array L, Array R>
requires (
    Numeric<typename L::value_type> &&
    Numeric<typename R::value_type>
)
L& operator-=(L& lhs, const R& rhs)
{
    return compound_op(lhs, rhs, std::minus<>{});
}

template<Array L, Array R>
requires (
    Numeric<typename L::value_type> &&
    Numeric<typename R::value_type>
)
L& operator*=(L& lhs, const R& rhs)
{
    return compound_op(lhs, rhs, std::multiplies<>{});
}

template<Array L, Array R>
requires (
    Numeric<typename L::value_type> &&
    Numeric<typename R::value_type>
)
L& operator/=(L& lhs, const R& rhs)
{
    return compound_op(lhs, rhs, std::divides<>{}, true);
}

template<Array A, Numeric S>
requires Numeric<typename A::value_type>
A& operator+=(A& lhs, const S& rhs)
{
    return compound_scalar_op(lhs, rhs, std::plus<>{});
}

template<Array A, Numeric S>
requires Numeric<typename A::value_type>
A& operator-=(A& lhs, const S& rhs)
{
    return compound_scalar_op(lhs, rhs, std::minus<>{});
}

template<Array A, Numeric S>
requires Numeric<typename A::value_type>
A& operator*=(A& lhs, const S& rhs)
{
    return compound_scalar_op(lhs, rhs, std::multiplies<>{});
}

template<Array A, Numeric S>
requires Numeric<typename A::value_type>
A& operator/=(A& lhs, const S& rhs)
{
    return compound_scalar_op(lhs, rhs, std::divides<>{}, true);
}
 
template<Array A>
requires Numeric<typename A::value_type>
A operator-(const A& arr)
{
    return arr * typename A::value_type{-1};
}
 
template<Array A>
requires Numeric<typename A::value_type>
A operator+(const A& arr)
{
    return arr;
}