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

// vector_operations.hpp
//
// Vector-operations implementation
//
// This file defines the vector-operation functions of the MKL algebra library.
// They provide arithmetic, math functions, and utilities for vectors.
//
// Main features:
// - Vector arithmetic (addition, subtraction, scalar multiplication, ...)
// - Inner product, outer product, and norm computation
// - Element-wise math functions (absolute value, square root, ...)
// - Implementation switching via computation policy

#ifndef SANGI_VECTOR_OPERATIONS_HPP
#define SANGI_VECTOR_OPERATIONS_HPP

#include "math/core/common.hpp"
#include "math/core/traits.hpp"
#include "math/concepts/algebraic_concepts.hpp"
#include "computation_policy.hpp"

#include <algorithm>
#include <cmath>
#include <functional>
#include <numeric>
#include <type_traits>

namespace sangi {
    namespace detail {
        namespace vector_operations {

            //-----------------------------------------------------------------------------
            // Basic vector operations
            //-----------------------------------------------------------------------------

            // Vector addition: result = a + b
            template<typename VectorA, typename VectorB, typename VectorResult,
                typename ComputationPolicy = detail::DefaultComputationPolicy<typename VectorA::value_type>>
                requires concepts::AdditiveGroup<typename VectorA::value_type>&&
            concepts::AdditiveGroup<typename VectorB::value_type>&&
                concepts::AdditiveGroup<typename VectorResult::value_type>
                void add(const VectorA& a, const VectorB& b, VectorResult& result) {
                ComputationPolicy::add(a, b, result);
            }

            // Vector subtraction: result = a - b
            template<typename VectorA, typename VectorB, typename VectorResult,
                typename ComputationPolicy = detail::DefaultComputationPolicy<typename VectorA::value_type>>
                requires concepts::AdditiveGroup<typename VectorA::value_type>&&
            concepts::AdditiveGroup<typename VectorB::value_type>&&
                concepts::AdditiveGroup<typename VectorResult::value_type>
                void subtract(const VectorA& a, const VectorB& b, VectorResult& result) {
                ComputationPolicy::subtract(a, b, result);
            }

            // Scalar multiplication: result = alpha * a
            template<typename VectorA, typename Scalar, typename VectorResult,
                typename ComputationPolicy = detail::DefaultComputationPolicy<typename VectorA::value_type>>
                requires concepts::Ring<typename VectorA::value_type>&&
            concepts::Ring<Scalar>&&
                concepts::Ring<typename VectorResult::value_type>
                void scale(const VectorA& a, const Scalar& alpha, VectorResult& result) {
                // Size check
                if (result.size() != a.size()) {
                    throw DimensionError("Vector scaling: size mismatch");
                }

                // First copy a into the result vector
                if (&result != &a) {
                    for (std::size_t i = 0; i < a.size(); ++i) {
                        result[i] = a[i];
                    }
                }

                // Then apply scaling
                ComputationPolicy::scale(result, alpha);
            }

            // Vector negation: result = -a
            template<typename VectorA, typename VectorResult,
                typename ComputationPolicy = detail::DefaultComputationPolicy<typename VectorA::value_type>>
                requires concepts::AdditiveGroup<typename VectorA::value_type>&&
            concepts::AdditiveGroup<typename VectorResult::value_type>
                void negate(const VectorA& a, VectorResult& result) {
                scale(a, typename VectorA::value_type(-1), result);
            }

            // Inner product: return = a . b
            template<typename VectorA, typename VectorB,
                typename ComputationPolicy = detail::DefaultComputationPolicy<typename VectorA::value_type>>
                requires concepts::Ring<typename VectorA::value_type>&&
            concepts::Ring<typename VectorB::value_type>
                auto dot_product(const VectorA& a, const VectorB& b)
                -> decltype(a[0] * b[0]) {
                return ComputationPolicy::dot_product(a, b);
            }

            // Element-wise product: result = a .* b (element-wise multiplication)
            template<typename VectorA, typename VectorB, typename VectorResult,
                typename ComputationPolicy = detail::DefaultComputationPolicy<typename VectorA::value_type>>
                requires concepts::Ring<typename VectorA::value_type>&&
            concepts::Ring<typename VectorB::value_type>&&
                concepts::Ring<typename VectorResult::value_type>
                void element_wise_multiply(const VectorA& a, const VectorB& b, VectorResult& result) {
                ComputationPolicy::element_wise_multiply(a, b, result);
            }

            // Element-wise division: result = a ./ b (element-wise division)
            template<typename VectorA, typename VectorB, typename VectorResult>
                requires concepts::Field<typename VectorA::value_type>&&
            concepts::Field<typename VectorB::value_type>&&
                concepts::Field<typename VectorResult::value_type>
                void element_wise_divide(const VectorA& a, const VectorB& b, VectorResult& result) {
                const auto size = a.size();

                // Size check
                if (b.size() != size || result.size() != size) {
                    throw DimensionError("Vector element-wise division: size mismatch");
                }

                // Element-wise division (with division-by-zero check)
                for (std::size_t i = 0; i < size; ++i) {
                    if (approximately_equal(b[i], numeric_traits<typename VectorB::value_type>::zero())) {
                        throw MathError("Vector element-wise division: division by zero");
                    }
                    result[i] = a[i] / b[i];
                }
            }

            // axpy operation: y += alpha * x
            template<typename VectorY, typename Scalar, typename VectorX,
                typename ComputationPolicy = detail::DefaultComputationPolicy<typename VectorY::value_type>>
                requires concepts::Ring<typename VectorY::value_type>&&
            concepts::Ring<Scalar>&&
                concepts::Ring<typename VectorX::value_type>
                void axpy(VectorY& y, const Scalar& alpha, const VectorX& x) {
                ComputationPolicy::axpy(y, alpha, x);
            }

            //-----------------------------------------------------------------------------
            // Math functions
            //-----------------------------------------------------------------------------

            // Apply a function to each element of a vector: result = f(a)
            template<typename VectorA, typename VectorResult, typename Function>
            void apply_function(const VectorA& a, VectorResult& result, Function f) {
                const auto size = a.size();

                // Size check
                if (result.size() != size) {
                    throw DimensionError("Vector function application: size mismatch");
                }

                // Apply the function to each element
                for (std::size_t i = 0; i < size; ++i) {
                    result[i] = f(a[i]);
                }
            }

            // Element-wise absolute value: result = |a|
            template<typename VectorA, typename VectorResult>
                requires requires(typename VectorA::value_type x) { { std::abs(x) } -> std::convertible_to<typename VectorResult::value_type>; }
            void abs(const VectorA& a, VectorResult& result) {
                apply_function(a, result, [](const auto& x) { return std::abs(x); });
            }

            // Element-wise square root: result = sqrt(a)
            template<typename VectorA, typename VectorResult>
                requires requires(typename VectorA::value_type x) { { std::sqrt(x) } -> std::convertible_to<typename VectorResult::value_type>; }
            void sqrt(const VectorA& a, VectorResult& result) {
                apply_function(a, result, [](const auto& x) {
                    if (x < numeric_traits<decltype(x)>::zero()) {
                        throw MathError("Vector sqrt: negative input");
                    }
                    return std::sqrt(x);
                    });
            }

            // Element-wise square: result = a^2
            template<typename VectorA, typename VectorResult>
                requires concepts::Ring<typename VectorA::value_type>&&
            concepts::Ring<typename VectorResult::value_type>
                void square(const VectorA& a, VectorResult& result) {
                apply_function(a, result, [](const auto& x) { return x * x; });
            }

            // Element-wise exponential: result = e^a
            template<typename VectorA, typename VectorResult>
                requires requires(typename VectorA::value_type x) { { std::exp(x) } -> std::convertible_to<typename VectorResult::value_type>; }
            void exp(const VectorA& a, VectorResult& result) {
                apply_function(a, result, [](const auto& x) { return std::exp(x); });
            }

            // Element-wise logarithm: result = log(a)
            template<typename VectorA, typename VectorResult>
                requires requires(typename VectorA::value_type x) { { std::log(x) } -> std::convertible_to<typename VectorResult::value_type>; }
            void log(const VectorA& a, VectorResult& result) {
                apply_function(a, result, [](const auto& x) {
                    if (x <= numeric_traits<decltype(x)>::zero()) {
                        throw MathError("Vector log: non-positive input");
                    }
                    return std::log(x);
                    });
            }

            // Element-wise power: result = a^exponent
            template<typename VectorA, typename VectorResult, typename Exponent>
                requires requires(typename VectorA::value_type x, Exponent e) { { std::pow(x, e) } -> std::convertible_to<typename VectorResult::value_type>; }
            void pow(const VectorA& a, VectorResult& result, const Exponent& exponent) {
                apply_function(a, result, [&exponent](const auto& x) {
                    return std::pow(x, exponent);
                    });
            }

            //-----------------------------------------------------------------------------
            // Norms and distances
            //-----------------------------------------------------------------------------

            // Vector L1 norm (sum of absolute values): ||a||_1
            template<typename VectorA>
                requires requires(typename VectorA::value_type x) { { std::abs(x) } -> std::convertible_to<typename VectorA::value_type>; }
            auto l1_norm(const VectorA& a) -> typename VectorA::value_type {
                using value_type = typename VectorA::value_type;
                value_type sum = numeric_traits<value_type>::zero();

                for (std::size_t i = 0; i < a.size(); ++i) {
                    sum += std::abs(a[i]);
                }

                return sum;
            }

            // Vector L2 norm (Euclidean norm): ||a||_2
            template<typename VectorA,
                typename ComputationPolicy = detail::DefaultComputationPolicy<typename VectorA::value_type>>
                requires concepts::Ring<typename VectorA::value_type>&&
                requires(typename VectorA::value_type x) { { std::sqrt(x) } -> std::convertible_to<typename VectorA::value_type>; }
            auto l2_norm(const VectorA& a) -> typename VectorA::value_type {
                return ComputationPolicy::norm2(a);
            }

            // Vector L-infinity norm (max absolute value): ||a||_inf
            template<typename VectorA>
                requires requires(typename VectorA::value_type x) { { std::abs(x) } -> std::convertible_to<typename VectorA::value_type>; }
            auto linf_norm(const VectorA& a) -> typename VectorA::value_type {
                if (a.empty()) {
                    return numeric_traits<typename VectorA::value_type>::zero();
                }

                using value_type = typename VectorA::value_type;
                value_type max_abs = std::abs(a[0]);

                for (std::size_t i = 1; i < a.size(); ++i) {
                    max_abs = std::max(max_abs, std::abs(a[i]));
                }

                return max_abs;
            }

            // L2 distance between vectors: ||a - b||_2
            template<typename VectorA, typename VectorB>
                requires concepts::AdditiveGroup<typename VectorA::value_type>&&
            concepts::AdditiveGroup<typename VectorB::value_type>&&
                requires(typename VectorA::value_type x) { { std::sqrt(x) } -> std::convertible_to<decltype(std::sqrt(std::declval<typename VectorA::value_type>()))>; }
            auto l2_distance(const VectorA& a, const VectorB& b)
                -> decltype(std::sqrt(std::declval<typename VectorA::value_type>())) {
                const auto size = a.size();

                // Size check
                if (b.size() != size) {
                    throw DimensionError("Vector L2 distance: size mismatch");
                }

                using result_type = decltype(a[0] - b[0]);
                result_type sum_squares = numeric_traits<result_type>::zero();

                for (std::size_t i = 0; i < size; ++i) {
                    const result_type diff = a[i] - b[i];
                    sum_squares += diff * diff;
                }

                return std::sqrt(sum_squares);
            }

            //-----------------------------------------------------------------------------
            // Statistics functions
            //-----------------------------------------------------------------------------

            // Minimum element of a vector
            template<typename VectorA>
                requires concepts::Comparable<typename VectorA::value_type>
            auto min(const VectorA& a) -> typename VectorA::value_type {
                if (a.empty()) {
                    throw IndexError("Vector min: empty vector");
                }

                return *std::min_element(a.begin(), a.end());
            }

            // Maximum element of a vector
            template<typename VectorA>
                requires concepts::Comparable<typename VectorA::value_type>
            auto max(const VectorA& a) -> typename VectorA::value_type {
                if (a.empty()) {
                    throw IndexError("Vector max: empty vector");
                }

                return *std::max_element(a.begin(), a.end());
            }

            // Sum of the vector's elements
            template<typename VectorA>
                requires concepts::AdditiveGroup<typename VectorA::value_type>
            auto sum(const VectorA& a) -> typename VectorA::value_type {
                if (a.empty()) {
                    return numeric_traits<typename VectorA::value_type>::zero();
                }

                using value_type = typename VectorA::value_type;
                return std::accumulate(a.begin(), a.end(), numeric_traits<value_type>::zero());
            }

            // Product of the vector's elements
            template<typename VectorA>
                requires concepts::MultiplicativeMonoid<typename VectorA::value_type>
            auto product(const VectorA& a) -> typename VectorA::value_type {
                if (a.empty()) {
                    return numeric_traits<typename VectorA::value_type>::one();
                }

                using value_type = typename VectorA::value_type;
                return std::accumulate(a.begin(), a.end(), numeric_traits<value_type>::one(),
                    std::multiplies<value_type>());
            }

            // Arithmetic mean of a vector
            template<typename VectorA>
                requires concepts::Field<typename VectorA::value_type>
            auto mean(const VectorA& a) -> typename VectorA::value_type {
                if (a.empty()) {
                    throw IndexError("Vector mean: empty vector");
                }

                using value_type = typename VectorA::value_type;
                const value_type total = sum(a);
                return total / static_cast<value_type>(a.size());
            }

            // Variance of a vector
            template<typename VectorA>
                requires concepts::Field<typename VectorA::value_type>
            auto variance(const VectorA& a) -> typename VectorA::value_type {
                if (a.empty() || a.size() == 1) {
                    throw IndexError("Vector variance: insufficient data");
                }

                const auto avg = mean(a);
                using value_type = typename VectorA::value_type;
                value_type sum_sq_diff = numeric_traits<value_type>::zero();

                for (std::size_t i = 0; i < a.size(); ++i) {
                    const value_type diff = a[i] - avg;
                    sum_sq_diff += diff * diff;
                }

                return sum_sq_diff / static_cast<value_type>(a.size() - 1);
            }

            // Standard deviation of a vector
            template<typename VectorA>
                requires concepts::Field<typename VectorA::value_type>&&
                requires(typename VectorA::value_type x) { { std::sqrt(x) } -> std::convertible_to<typename VectorA::value_type>; }
            auto standard_deviation(const VectorA& a) -> typename VectorA::value_type {
                return std::sqrt(variance(a));
            }

            // Inner product of vectors with numeric-state awareness
            template<typename VectorA, typename VectorB>
                requires HasNumericState<typename VectorA::value_type>&&
            HasNumericState<typename VectorB::value_type>
                auto dot_product_with_state(const VectorA& a, const VectorB& b)
                -> typename VectorA::value_type {
                const auto size = a.size();

                // Size check
                if (b.size() != size) {
                    throw DimensionError("Vector dot product with state: size mismatch");
                }

                // Check for special states
                for (std::size_t i = 0; i < size; ++i) {
                    if (numeric_state_traits<typename VectorA::value_type>::isNaN(a[i]) ||
                        numeric_state_traits<typename VectorB::value_type>::isNaN(b[i])) {
                        // If any value is NaN, return NaN
                        return numeric_traits<typename VectorA::value_type>::quiet_NaN();
                    }
                    if (numeric_state_traits<typename VectorA::value_type>::isInfinite(a[i]) ||
                        numeric_state_traits<typename VectorB::value_type>::isInfinite(b[i])) {
                        // Handle infinity case explicitly
                        // e.g. inf * 0 = NaN, inf * non-zero = inf
                        if ((numeric_state_traits<typename VectorA::value_type>::isInfinite(a[i]) &&
                            b[i] == numeric_traits<typename VectorB::value_type>::zero()) ||
                            (numeric_state_traits<typename VectorB::value_type>::isInfinite(b[i]) &&
                                a[i] == numeric_traits<typename VectorA::value_type>::zero())) {
                            return numeric_traits<typename VectorA::value_type>::quiet_NaN();
                        }
                        // For any other infinity, return infinity
                        // (the sign depends on the actual result)
                        return numeric_traits<typename VectorA::value_type>::infinity();
                    }
                }

                // Ordinary inner-product computation
                using result_type = typename VectorA::value_type;
                result_type sum = numeric_traits<result_type>::zero();

                for (std::size_t i = 0; i < size; ++i) {
                    sum += a[i] * b[i];
                }

                return sum;
            }

            // L2 norm of a vector with numeric-state awareness
            template<typename VectorA>
                requires HasNumericState<typename VectorA::value_type>
            auto l2_norm_with_state(const VectorA& a) -> typename VectorA::value_type {
                using value_type = typename VectorA::value_type;

                // Check for special states
                for (std::size_t i = 0; i < a.size(); ++i) {
                    if (numeric_state_traits<value_type>::isNaN(a[i])) {
                        // If a NaN is present, return NaN
                        return numeric_traits<value_type>::quiet_NaN();
                    }
                    if (numeric_state_traits<value_type>::isInfinite(a[i])) {
                        // If any element is infinite, return infinity
                        return numeric_traits<value_type>::infinity();
                    }
                }

                // Ordinary norm computation
                return l2_norm(a);
            }

} // namespace vector_operations
} // namespace detail
} // namespace sangi

#endif // SANGI_VECTOR_OPERATIONS_HPP