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

// ModularIntTraits.hpp
#ifndef SANGI_MODULAR_INT_TRAITS_HPP
#define SANGI_MODULAR_INT_TRAITS_HPP
// Integration of modular integer class with algebraic traits (C++20/23 concept-aware)
//
// This file provides customization points to integrate the ModularInt<P> class
// with the sangi library's algebraic traits. This allows ModularInt to
// cooperate seamlessly with other library components such as
// matrices, vectors, and numerical algorithms.
//
// Main features:
// - Integration with C++20/23 concepts
// - Customization of additive/multiplicative identity for ModularInt
// - Integration with the sangi::concepts namespace
// - Integration with MKL BLAS traits (e.g. sparse-matrix operations)
// 
// Usage examples:
// - Matrix operations using ModularInt
// - Linear algebra over finite fields
// - Implementation of number-theoretic algorithms

#include "ModularInt.hpp"
#include "CRT.hpp"
#include <math/concepts/algebraic_concepts.hpp>
#include <math/core/common.hpp>

#include <vector>
#include <cstdint>
#include <numeric>
#include <optional>

namespace sangi {

// Specializations to make ModularInt fit the concepts
namespace concepts {

// Mark that ModularInt satisfies the Field concept
template <int P>
inline constexpr bool is_field_v<ModularInt<P>> = true;

// Note: Z/pZ is not an ordered field, so is_ordered_field_v is false
template <int P>
inline constexpr bool is_ordered_field_v<ModularInt<P>> = false;

// Mark that ModularInt satisfies the MklOptimizedType concept
template <int P>
inline constexpr bool is_mkl_optimized_v<ModularInt<P>> = false;  // MKL does not directly support it

// Specializations of operation traits
template <int P>
struct field_operations<ModularInt<P>> {
    static constexpr ModularInt<P> add(const ModularInt<P>& a, const ModularInt<P>& b) noexcept {
        return a + b;
    }

    static constexpr ModularInt<P> subtract(const ModularInt<P>& a, const ModularInt<P>& b) noexcept {
        return a - b;
    }

    static constexpr ModularInt<P> multiply(const ModularInt<P>& a, const ModularInt<P>& b) noexcept {
        return a * b;
    }

    static ModularInt<P> divide(const ModularInt<P>& a, const ModularInt<P>& b) {
        return a / b;
    }

    static constexpr ModularInt<P> negate(const ModularInt<P>& a) noexcept {
        return -a;
    }

    static constexpr ModularInt<P> zero() noexcept {
        return ModularInt<P>(0);
    }

    static constexpr ModularInt<P> one() noexcept {
        return ModularInt<P>(1);
    }

    static constexpr bool equals(const ModularInt<P>& a, const ModularInt<P>& b) noexcept {
        return a == b;
    }

    static constexpr bool less_than(const ModularInt<P>& a, const ModularInt<P>& b) noexcept {
        return a < b;
    }
};

} // namespace concepts

// MKL support
#if SANGI_HAS_MKL
namespace detail {
namespace matrix_mkl {

// MKL operation interface for ModularInt (MKL has no direct support, so a custom implementation)
template <int P>
class ModularIntMklOperations {
public:
    // Vector addition: y = a*x + y
    static void axpy(int n, ModularInt<P> alpha, const ModularInt<P>* x, int incx, 
                    ModularInt<P>* y, int incy) {
        for (int i = 0; i < n; ++i) {
            y[i * incy] += alpha * x[i * incx];
        }
    }

    // Vector dot product: result = x . y
    static ModularInt<P> dot(int n, const ModularInt<P>* x, int incx, 
                           const ModularInt<P>* y, int incy) {
        ModularInt<P> result(0);
        for (int i = 0; i < n; ++i) {
            result += x[i * incx] * y[i * incy];
        }
        return result;
    }

    // Matrix multiplication: C = alpha*A*B + beta*C
    static void gemm(int layout, int transA, int transB, int m, int n, int k,
                    ModularInt<P> alpha, const ModularInt<P>* A, int lda,
                    const ModularInt<P>* B, int ldb, ModularInt<P> beta,
                    ModularInt<P>* C, int ldc) {
        
        // Apply beta to C
        if (beta != ModularInt<P>(1)) {
            for (int i = 0; i < m; ++i) {
                for (int j = 0; j < n; ++j) {
                    C[i*ldc + j] = beta * C[i*ldc + j];
                }
            }
        }
        
        // Only row-major format is supported (simplified implementation)
        if (layout != CblasRowMajor || transA != CblasNoTrans || transB != CblasNoTrans) {
            throw MathError("Unsupported matrix layout or transposition for ModularInt");
        }

        // Compute A*B and add to C
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                ModularInt<P> sum(0);
                for (int l = 0; l < k; ++l) {
                    sum += A[i*lda + l] * B[l*ldb + j];
                }
                C[i*ldc + j] += alpha * sum;
            }
        }
    }

    // Implement additional MKL-compatible functions as needed
};

} // namespace matrix_mkl
} // namespace detail
#endif // SANGI_HAS_MKL

// Number-theoretic functions specific to ModularInt
// extended_gcd, chinese_remainder_theorem, and CRT class are defined in CRT.hpp

/**
 * @brief Fast modular exponentiation
 * @tparam P Modulus
 * @param base Base
 * @param exponent Exponent
 * @return Computation result
 */
template <int P>
constexpr ModularInt<P> mod_pow(const ModularInt<P>& base, int64_t exponent) {
    return base.pow(exponent);
}

/**
 * @brief Compute a primitive root
 * @tparam P Prime modulus
 * @return A primitive root of P, or nullopt if none is found
 */
template <int P>
std::optional<ModularInt<P>> find_primitive_root() {
    if (P <= 1) {
        return std::nullopt;
    }

    // Prime factorization of P-1
    int64_t phi = P - 1;
    std::vector<int64_t> prime_factors;
    
    // Prime factorization (naive algorithm)
    int64_t n = phi;
    for (int64_t i = 2; i * i <= n; ++i) {
        if (n % i == 0) {
            prime_factors.push_back(i);
            while (n % i == 0) {
                n /= i;
            }
        }
    }
    if (n > 1) {
        prime_factors.push_back(n);
    }

    // Search for primitive-root candidates
    for (int g = 2; g < P; ++g) {
        bool is_primitive = true;
        ModularInt<P> candidate(g);
        
        for (int64_t prime : prime_factors) {
            if (candidate.pow(phi / prime) == ModularInt<P>(1)) {
                is_primitive = false;
                break;
            }
        }
        
        if (is_primitive) {
            return ModularInt<P>(g);
        }
    }
    
    return std::nullopt;
}

/**
 * @brief Fast modular exponentiation (non-member function)
 * @tparam P Modulus
 * @param base Base
 * @param exponent Exponent
 * @return Computation result
 */
template <int P>
constexpr ModularInt<P> pow_mod(int64_t base, int64_t exponent) {
    return ModularInt<P>(base).pow(exponent);
}

/**
 * @brief Modular inverse (using Fermat's little theorem)
 * @tparam P Prime modulus
 * @param a The value whose inverse is sought
 * @return Modular inverse of a
 */
template <int P>
constexpr ModularInt<P> mod_inverse(int64_t a) {
    return ModularInt<P>(a).inverse();
}

/**
 * @brief Modular congruence solver
 * @tparam P Modulus
 * @param a Coefficient
 * @param b Constant term
 * @return Array of solutions x to ax = b (mod P)
 */
template <int P>
std::vector<ModularInt<P>> solve_congruence(int64_t a, int64_t b) {
    a = (a % P + P) % P;
    b = (b % P + P) % P;
    
    std::vector<ModularInt<P>> solutions;
    
    // Compute the GCD of a and P
    int64_t g = std::gcd(a, static_cast<int64_t>(P));
    
    if (b % g != 0) {
        // No solution
        return solutions;
    }
    
    // a' x = b' (mod P'), where a' = a/g, b' = b/g, P' = P/g
    a /= g;
    b /= g;
    int64_t mod_prime = P / g;
    
    // Compute the inverse of a'
    ModularInt<P> a_inv = ModularInt<P>(a).inverse();
    
    // Base solution x_0 = a^(-1) * b (mod P/g)
    ModularInt<P> x0 = a_inv * ModularInt<P>(b);
    
    // Generate all solutions: x_k = x_0 + k * (P/g) for k = 0, 1, ..., g-1
    for (int64_t k = 0; k < g; ++k) {
        solutions.push_back(x0 + ModularInt<P>(k * mod_prime));
    }
    
    return solutions;
}

} // namespace sangi

#endif // SANGI_MODULAR_INT_TRAITS_HPP