#ifndef ALGOIM_UVECTOR_HPP
#define ALGOIM_UVECTOR_HPP

// algoim::uvector implements a vector with compile-time known extent, including
// arithmetic expressions and various utility methods. It could be viewed as, e.g.,
// a light-weight, self-contained, re-implementation of blitz's TinyVector<T,N>,
// (https://en.wikipedia.org/wiki/Blitz%2B%2B), using more modern C++17/20
// techniques.

#include <type_traits>
#include <algorithm>
#include <cstring>
#include <iostream>
#include <cmath>
#include "real.hpp"

namespace algoim
{
template <typename T, int N>
class uvector;

template <int N, typename T>
struct uvector_expr {
    T eval;
};

namespace detail
{
template <typename T, int N>
constexpr auto eval(const uvector<T, N>& e, int i)
{
    return e(i);
}

template <int N, typename T>
constexpr auto eval(const uvector_expr<N, T>& expr, int i)
{
    return expr.eval(i);
}

template <typename E>
constexpr auto eval(const E& e, int i)
{
    return e;
}

template <int N, typename T>
constexpr auto make_uvector_expr(const T& eval)
{
    return uvector_expr<N, T>{eval};
}

template <typename T>
struct extent_of {
    static constexpr int value = 0;
};

template <typename T, int N>
struct extent_of<uvector<T, N>> {
    static constexpr int value = N;
};

template <int N, typename T>
struct extent_of<uvector_expr<N, T>> {
    static constexpr int value = N;
};

template <typename T>
using extent = std::integral_constant<int, extent_of<std::decay_t<T>>::value>;
} // namespace detail

// binary uvector expressions
#define algoim_decl_binary_op(op)                                                                                              \
    template <typename E1, typename E2,                                                                                        \
              std::enable_if_t<(detail::extent<E1>::value > 0) || (detail::extent<E2>::value > 0), bool> = true>               \
    auto operator op(E1&& e1, E2&& e2)                                                                                         \
    {                                                                                                                          \
        static_assert(detail::extent<E1>::value == 0 || detail::extent<E2>::value == 0                                         \
                          || detail::extent<E1>::value == detail::extent<E2>::value,                                           \
                      "incompatible extents in algoim::uvector expression");                                                   \
        constexpr int N = std::max(detail::extent<E1>::value, detail::extent<E2>::value);                                      \
        if constexpr (std::is_lvalue_reference<E1&&>::value)                                                                   \
            if constexpr (std::is_lvalue_reference<E2&&>::value)                                                               \
                return detail::make_uvector_expr<N>([&e1, &e2](int i) { return detail::eval(e1, i) op detail::eval(e2, i); }); \
            else                                                                                                               \
                return detail::make_uvector_expr<N>([&e1, e2](int i) { return detail::eval(e1, i) op detail::eval(e2, i); });  \
        else if constexpr (std::is_lvalue_reference<E2&&>::value)                                                              \
            return detail::make_uvector_expr<N>([e1, &e2](int i) { return detail::eval(e1, i) op detail::eval(e2, i); });      \
        else                                                                                                                   \
            return detail::make_uvector_expr<N>([e1, e2](int i) { return detail::eval(e1, i) op detail::eval(e2, i); });       \
    }
algoim_decl_binary_op(+) algoim_decl_binary_op(-) algoim_decl_binary_op(*) algoim_decl_binary_op(/) algoim_decl_binary_op(<)
    algoim_decl_binary_op(<=) algoim_decl_binary_op(>) algoim_decl_binary_op(>=) algoim_decl_binary_op(==)
        algoim_decl_binary_op(!=) algoim_decl_binary_op(%)
#undef algoim_decl_binary_op

    // unary negation operator for uvector or uvector expressions
    template <typename E, std::enable_if_t<detail::extent<E>::value >= 1, bool> = true>
    auto operator-(const E& e)
{
    constexpr int N = detail::extent<E>::value;
    if constexpr (std::is_lvalue_reference<E&&>::value)
        return detail::make_uvector_expr<N>([&e](int i) { return -detail::eval(e, i); });
    else
        return detail::make_uvector_expr<N>([e](int i) { return -detail::eval(e, i); });
}

// abs of uvector or uvector expressions
template <typename E, std::enable_if_t<detail::extent<E>::value >= 1, bool> = true>
auto abs(const E& e)
{
    using std::abs;
    constexpr int N = detail::extent<E>::value;
    if constexpr (std::is_lvalue_reference<E&&>::value)
        return detail::make_uvector_expr<N>([&e](int i) { return abs(detail::eval(e, i)); });
    else
        return detail::make_uvector_expr<N>([e](int i) { return abs(detail::eval(e, i)); });
}

// uvector class
template <typename T, int N>
class uvector
{
    T data_[N];

public:
    // default constructor zero initialises
    uvector() { memset(data_, 0, sizeof(data_)); }

    template <typename... U, std::enable_if_t<N >= 2 && sizeof...(U) == N, bool> = true>
    explicit uvector(U&&... args) : data_{static_cast<T>(std::forward<U>(args))...}
    {
    }

    template <typename U>
    uvector(const U& x)
    {
        *this = x;
    }

    T& operator()(int i) { return data_[i]; }

    const T& operator()(int i) const { return data_[i]; }

    T* data() { return data_; }

    const T* data() const { return data_; }

// assignment and update operators, e.g., operator=, operator+=, operator*=, etc.
#define algoim_decl_assign_op(op)                                                     \
    template <typename E>                                                             \
    uvector& operator op(const E & x)                                                 \
    {                                                                                 \
        static_assert(detail::extent<E>::value == 0 || detail::extent<E>::value == N, \
                      "incompatible extents in algoim::uvector expression");          \
        for (int i = 0; i < N; ++i) data_[i] op detail::eval(x, i);                   \
        return *this;                                                                 \
    }
    algoim_decl_assign_op(=) algoim_decl_assign_op(+=) algoim_decl_assign_op(-=) algoim_decl_assign_op(*=)
        algoim_decl_assign_op(/=)
#undef algoim_decl_assign_op
};

using uvector3f = uvector<real, 3>;
using uvector2f = uvector<real, 2>;
using uvector3i = uvector<int, 3>;
using uvector2i = uvector<int, 2>;
using uvector3  = uvector3f;
using uvector2  = uvector2f;

// some methods, particularly those which remove components of a vector, benefit
// from having a zero-length uvector; it is simply an empty class
template <typename T>
class uvector<T, 0>
{
};

template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
std::ostream& operator<<(std::ostream& o, const T& u)
{
    constexpr int N = detail::extent<T>::value;
    o << '(';
    for (int i = 0; i < N; ++i) o << detail::eval(u, i) << (i + 1 < N ? ',' : ')');
    return o;
}

// minimum component of a uvector or uvector expression
template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
auto min(const T& u)
{
    constexpr int N = detail::extent<T>::value;
    auto          x = detail::eval(u, 0);
    for (int i = 1; i < N; ++i) {
        auto y = detail::eval(u, i);
        if (y < x) x = y;
    }
    return x;
}

// maximum component of a uvector or uvector expression
template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
auto max(const T& u)
{
    constexpr int N = detail::extent<T>::value;
    auto          x = detail::eval(u, 0);
    for (int i = 1; i < N; ++i) {
        auto y = detail::eval(u, i);
        if (y > x) x = y;
    }
    return x;
}

// index where the minimum of a uvector or uvector expression is attained; if the minimum
// is attained at multiple places, the smallest corresponding index is returned
template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
int argmin(const T& u)
{
    constexpr int N   = detail::extent<T>::value;
    auto          x   = detail::eval(u, 0);
    int           ind = 0;
    for (int i = 1; i < N; ++i) {
        auto y = detail::eval(u, i);
        if (y < x) {
            x   = y;
            ind = i;
        }
    }
    return ind;
}

// index where the maximum of a uvector or uvector expression is attained; if the maximum
// is attained at multiple places, the smallest corresponding index is returned
template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
int argmax(const T& u)
{
    constexpr int N   = detail::extent<T>::value;
    auto          x   = detail::eval(u, 0);
    int           ind = 0;
    for (int i = 1; i < N; ++i) {
        auto y = detail::eval(u, i);
        if (y > x) {
            x   = y;
            ind = i;
        }
    }
    return ind;
}

// sum of components of a uvector or uvector expression
template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
auto sum(const T& u)
{
    constexpr int N = detail::extent<T>::value;
    auto          x = detail::eval(u, 0);
    for (int i = 1; i < N; ++i) x += detail::eval(u, i);
    return x;
}

// product of components of a uvector or uvector expression
template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
auto prod(const T& u)
{
    constexpr int N = detail::extent<T>::value;
    auto          x = detail::eval(u, 0);
    for (int i = 1; i < N; ++i) x *= detail::eval(u, i);
    return x;
}

// product of components of a uvector or uvector expression, excluding
// a given index; the empty product returns one
template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
auto prod(const T& u, int excluded_index)
{
    constexpr int N = detail::extent<T>::value;
    auto          x = detail::eval(u, 0); // deduce return type
    x               = 1;
    for (int i = 0; i < N; ++i)
        if (i != excluded_index) x *= detail::eval(u, i);
    return x;
}

// dot product of two uvector or uvector expressions
template <
    typename E1,
    typename E2,
    std::enable_if_t<detail::extent<E1>::value == detail::extent<E2>::value && detail::extent<E1>::value >= 1, bool> = true>
auto dot(const E1& u, const E2& v)
{
    constexpr int N = detail::extent<E1>::value;
    auto          x = detail::eval(u, 0) * detail::eval(v, 0);
    for (int i = 1; i < N; ++i) x += detail::eval(u, i) * detail::eval(v, i);
    return x;
}

// cross product of two uvector3 or uvector3 expressions
template <typename E1,
          typename E2,
          std::enable_if_t<detail::extent<E1>::value == 3 && detail::extent<E2>::value == 3, bool> = true>

auto cross(const E1& u, const E2& v)
{
    return uvector<real, 3>(detail::eval(u, 1) * detail::eval(v, 2) - detail::eval(u, 2) * detail::eval(v, 1),
                            detail::eval(u, 2) * detail::eval(v, 0) - detail::eval(u, 0) * detail::eval(v, 2),
                            detail::eval(u, 0) * detail::eval(v, 1) - detail::eval(u, 1) * detail::eval(v, 0));
}

// component-wise max of two uvectors
template <typename T, int N>
auto max(const uvector<T, N>& u, const uvector<T, N>& v)
{
    using std::max;
    uvector<T, N> x;
    for (int i = 0; i < N; ++i) x(i) = max(u(i), v(i));
    return x;
}

// returns false iff all components of u yield false when typecast to a bool
template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
bool any(const T& u)
{
    constexpr int N = detail::extent<T>::value;
    bool          x = false;
    for (int i = 0; i < N; ++i) x |= static_cast<bool>(detail::eval(u, i));
    return x;
}

// returns true iff all components of u yield true when typecast to a bool
template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
bool all(const T& u)
{
    constexpr int N = detail::extent<T>::value;
    bool          x = true;
    for (int i = 0; i < N; ++i) x &= static_cast<bool>(detail::eval(u, i));
    return x;
}

// squared Euclidean norm of a uvector or uvector expression
template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
auto sqrnorm(const T& u)
{
    return dot(u, u);
}

// Euclidean norm of a uvector or uvector expression
template <typename T, std::enable_if_t<detail::extent<T>::value >= 1, bool> = true>
auto norm(const T& u)
{
    using std::sqrt;
    return sqrt(sqrnorm(u));
}

// alter the extent of a uvector by truncation or zero backfill
template <int M, typename T, int N>
uvector<T, M> alter_extent(const uvector<T, N>& u)
{
    uvector<T, M> v;
    for (int i = 0; i < (M < N ? M : N); ++i) // truncate if M < N
        v(i) = u(i);
    for (int i = N; i < M; ++i)               // backfill with zeros if M > N
        memset(&v(i), 0, sizeof(v(i)));
    return v;
}

// remove a component from a uvector
template <typename T, int N>
uvector<T, N - 1> remove_component(const uvector<T, N>& u, int index)
{
    uvector<T, N - 1> v;
    if constexpr (N > 1)
        for (int i = 0; i < N - 1; ++i) v(i) = u(i < index ? i : i + 1);
    return v;
}

// add a component to a vector, consistent with remove_component semantics
template <typename T, int N>
uvector<T, N + 1> add_component(const uvector<T, N>& u, int index, T value)
{
    if constexpr (N == 0)
        return uvector<T, 1>(value);
    else {
        uvector<T, N + 1> v;
        for (int i = 0; i < N + 1; ++i) v(i) = i < index ? u(i) : (i == index ? value : u(i - 1));
        return v;
    }
}

// set the i'th component of a uvector to a given value
template <typename T, int N>
uvector<T, N> set_component(uvector<T, N> u, int i, T value)
{
    u(i) = value;
    return u;
}

// increment the i'th component of a uvector by a given value
template <typename T, int N>
uvector<T, N> inc_component(uvector<T, N> u, int i, T value)
{
    u(i) += value;
    return u;
}
} // namespace algoim

#endif