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

// ModularIntTraits.hpp
#ifndef CALX_MODULAR_INT_TRAITS_HPP
#define CALX_MODULAR_INT_TRAITS_HPP
// モジュラー整数クラスと代数系トレイトの統合（C++20/23コンセプト対応版）
//
// このファイルは、ModularInt<P>クラスをcalxライブラリの代数系トレイトと
// 統合するためのカスタマイゼーションポイントを提供します。これにより、ModularIntを
// 行列、ベクトル、数値計算アルゴリズムなど、ライブラリの他のコンポーネントと
// シームレスに連携させることが可能になります。
//
// 主な機能:
// - C++20/23コンセプトとの統合
// - ModularInt用の加法・乗法単位元のカスタマイズ
// - calx::concepts名前空間との統合
// - MKL BLASトレイトとの統合（疎行列演算など）
// 
// 使用例:
// - ModularIntを使った行列演算
// - 有限体上の線形代数計算
// - 数論アルゴリズムの実装

#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 calx {

// ModularIntをコンセプトに適合させるための特殊化
namespace concepts {

// ModularIntがField（体）コンセプトを満たすことを明示
template <int P>
inline constexpr bool is_field_v<ModularInt<P>> = true;

// 注意: 有限体 Z/pZ は順序体ではないため is_ordered_field_v は false
template <int P>
inline constexpr bool is_ordered_field_v<ModularInt<P>> = false;

// ModularIntがMklOptimizedType（MKL最適化型）コンセプトを満たすことを明示
template <int P>
inline constexpr bool is_mkl_optimized_v<ModularInt<P>> = false;  // MKLは直接対応していない

// 演算特性の特殊化
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対応
#if CALX_HAS_MKL
namespace detail {
namespace matrix_mkl {

// ModularInt用のMKL演算インターフェース（MKLは対応していないのでカスタム実装）
template <int P>
class ModularIntMklOperations {
public:
    // ベクトル加算: 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];
        }
    }

    // ベクトルのドット積: 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;
    }

    // 行列乗算: 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) {
        
        // Cに対するbetaの適用
        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];
                }
            }
        }
        
        // 行優先形式のみサポート（実装を簡略化）
        if (layout != CblasRowMajor || transA != CblasNoTrans || transB != CblasNoTrans) {
            throw MathError("Unsupported matrix layout or transposition for ModularInt");
        }

        // A*Bの計算と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;
            }
        }
    }

    // 追加のMKL互換関数を必要に応じて実装
};

} // namespace matrix_mkl
} // namespace detail
#endif // CALX_HAS_MKL

// ModularInt特有の数論関数
// extended_gcd, chinese_remainder_theorem, CRT クラスは CRT.hpp に定義

/**
 * @brief 高速モジュラー累乗計算
 * @tparam P モジュラス
 * @param base 底
 * @param exponent 指数
 * @return 計算結果
 */
template <int P>
constexpr ModularInt<P> mod_pow(const ModularInt<P>& base, int64_t exponent) {
    return base.pow(exponent);
}

/**
 * @brief 原始根の計算
 * @tparam P 素数モジュラス
 * @return Pの原始根またはnullopt（見つからない場合）
 */
template <int P>
std::optional<ModularInt<P>> find_primitive_root() {
    if (P <= 1) {
        return std::nullopt;
    }

    // P-1の素因数分解
    int64_t phi = P - 1;
    std::vector<int64_t> prime_factors;
    
    // 素因数分解（単純アルゴリズム）
    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);
    }

    // 原始根の候補を検索
    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 高速モジュラーべき乗計算（非メンバー関数）
 * @tparam P モジュラス
 * @param base 底
 * @param exponent 指数
 * @return 計算結果
 */
template <int P>
constexpr ModularInt<P> pow_mod(int64_t base, int64_t exponent) {
    return ModularInt<P>(base).pow(exponent);
}

/**
 * @brief モジュラー逆元計算（フェルマーの小定理使用）
 * @tparam P 素数モジュラス
 * @param a 逆元を求める数
 * @return aのモジュラー逆元
 */
template <int P>
constexpr ModularInt<P> mod_inverse(int64_t a) {
    return ModularInt<P>(a).inverse();
}

/**
 * @brief モジュラー合同式ソルバー
 * @tparam P モジュラス
 * @param a 係数
 * @param b 定数項
 * @return ax ≡ b (mod P)の解xの配列
 */
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;
    
    // aとPの最大公約数を計算
    int64_t g = std::gcd(a, static_cast<int64_t>(P));
    
    if (b % g != 0) {
        // 解なし
        return solutions;
    }
    
    // a'x ≡ b' (mod P')、ここでa'=a/g, b'=b/g, P'=P/g
    a /= g;
    b /= g;
    int64_t mod_prime = P / g;
    
    // a'の逆元を計算
    ModularInt<P> a_inv = ModularInt<P>(a).inverse();
    
    // 基本解x_0 = a^(-1) * b (mod P/g)
    ModularInt<P> x0 = a_inv * ModularInt<P>(b);
    
    // 全ての解を生成: 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 calx

#endif // CALX_MODULAR_INT_TRAITS_HPP