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

// DirichletCharacter.hpp
//
// Dirichlet 指標 (Dirichlet Characters)
//
// (Z/qZ)* 上の群準同型 χ: (Z/qZ)* → C*
//
// 主な API:
//   DirichletCharacter(q, values) — 指標の定義
//   DirichletCharacter::principal(q) — 主指標
//   operator()(n)                 — χ(n) の評価
//   conductor()                   — 導手
//   isPrimitive()                 — 原始指標判定
//   order()                       — 位数
//   dirichletL(chi, s, terms)     — L(s, χ) の計算

#ifndef CALX_CORE_DIRICHLET_CHARACTER_HPP
#define CALX_CORE_DIRICHLET_CHARACTER_HPP

#include <cstdint>
#include <cmath>
#include <complex>
#include <stdexcept>
#include <vector>
#include <algorithm>
#include <numeric>
#include <numbers>

namespace calx {

/**
 * @brief Dirichlet 指標 χ mod q
 *
 * 値は根 of unity として std::complex<double> で保持。
 * gcd(n, q) > 1 のとき χ(n) = 0。
 */
class DirichletCharacter {
public:
    /// テーブルから構築 (values[n] = χ(n), 0 ≤ n < q)
    DirichletCharacter(int64_t q, const std::vector<std::complex<double>>& values)
        : q_(q), values_(values) {
        if (q < 1) throw std::invalid_argument("DirichletCharacter: q must be >= 1");
        values_.resize(static_cast<std::size_t>(q), 0.0);
    }

    /// 主指標 χ_0: gcd(n,q)=1 → 1, otherwise 0
    static DirichletCharacter principal(int64_t q) {
        std::vector<std::complex<double>> vals(static_cast<std::size_t>(q), 0.0);
        for (int64_t n = 0; n < q; ++n) {
            if (gcd(n, q) == 1)
                vals[static_cast<std::size_t>(n)] = 1.0;
        }
        return DirichletCharacter(q, vals);
    }

    /// Legendre 記号 (p/·) を指標として構築 (p は奇素数)
    static DirichletCharacter legendreCharacter(int64_t p) {
        if (p < 3) throw std::invalid_argument("legendreCharacter: p must be odd prime >= 3");
        std::vector<std::complex<double>> vals(static_cast<std::size_t>(p), 0.0);
        for (int64_t n = 1; n < p; ++n) {
            // Euler の規準: (n/p) = n^{(p-1)/2} mod p
            int64_t result = modpow(n % p, (p - 1) / 2, p);
            if (result == p - 1) result = -1;
            vals[static_cast<std::size_t>(n)] = static_cast<double>(result);
        }
        return DirichletCharacter(p, vals);
    }

    /// χ(n) の評価
    std::complex<double> operator()(int64_t n) const {
        n = ((n % q_) + q_) % q_;
        return values_[static_cast<std::size_t>(n)];
    }

    int64_t modulus() const { return q_; }

    /// 位数 (χ^k = χ_0 となる最小の k)
    int64_t order() const {
        // χ の値から位数を推定
        // すべての非ゼロ値が ζ_k (k 乗根) ならば位数 = k
        int64_t ord = 1;
        for (int64_t n = 0; n < q_; ++n) {
            auto v = values_[static_cast<std::size_t>(n)];
            if (std::abs(v) < 1e-10) continue;
            // v = e^{2πi·r/s} → s を求める
            double angle = std::arg(v);
            if (std::abs(angle) < 1e-10) continue;
            // angle = 2π · p/q → q を推定
            double frac = angle / (2.0 * std::numbers::pi);
            // 最良有理近似で q を求める
            for (int64_t d = 1; d <= 100; ++d) {
                double nd = frac * static_cast<double>(d);
                if (std::abs(nd - std::round(nd)) < 1e-8) {
                    ord = lcm(ord, d);
                    break;
                }
            }
        }
        return ord;
    }

    /// 導手 (conductor): χ が mod f で定義できる最小の f
    int64_t conductor() const {
        // q の約数 f を小さい順に試す
        for (int64_t f = 1; f <= q_; ++f) {
            if (q_ % f != 0) continue;
            // χ が mod f で well-defined かチェック
            bool valid = true;
            for (int64_t a = 0; a < q_ && valid; ++a) {
                if (gcd(a, q_) != 1) continue;
                for (int64_t b = a + f; b < q_ && valid; b += f) {
                    if (gcd(b, q_) != 1) continue;
                    if (std::abs(values_[static_cast<std::size_t>(a)] -
                                values_[static_cast<std::size_t>(b)]) > 1e-8)
                        valid = false;
                }
            }
            if (valid) return f;
        }
        return q_;
    }

    /// 原始指標判定 (conductor = q)
    bool isPrimitive() const {
        return conductor() == q_;
    }

    /// 共役指標 χ̄(n) = conj(χ(n))
    DirichletCharacter conjugate() const {
        std::vector<std::complex<double>> conj_vals(values_.size());
        for (std::size_t i = 0; i < values_.size(); ++i)
            conj_vals[i] = std::conj(values_[i]);
        return DirichletCharacter(q_, conj_vals);
    }

    /// 指標の積 (χ₁ · χ₂)(n) = χ₁(n) · χ₂(n)
    friend DirichletCharacter operator*(const DirichletCharacter& a,
                                         const DirichletCharacter& b) {
        if (a.q_ != b.q_)
            throw std::invalid_argument("DirichletCharacter: different moduli");
        std::vector<std::complex<double>> vals(static_cast<std::size_t>(a.q_));
        for (int64_t n = 0; n < a.q_; ++n) {
            vals[static_cast<std::size_t>(n)] =
                a.values_[static_cast<std::size_t>(n)] *
                b.values_[static_cast<std::size_t>(n)];
        }
        return DirichletCharacter(a.q_, vals);
    }

private:
    int64_t q_;
    std::vector<std::complex<double>> values_;

    static int64_t gcd(int64_t a, int64_t b) {
        a = std::abs(a); b = std::abs(b);
        while (b) { a %= b; std::swap(a, b); }
        return a;
    }

    static int64_t lcm(int64_t a, int64_t b) {
        return a / gcd(a, b) * b;
    }

    static int64_t modpow(int64_t base, int64_t exp, int64_t mod) {
        int64_t result = 1;
        base %= mod;
        while (exp > 0) {
            if (exp & 1) result = (result * base) % mod;
            base = (base * base) % mod;
            exp >>= 1;
        }
        return result;
    }
};

/**
 * @brief Dirichlet L 関数 L(s, χ) = Σ χ(n)/n^s
 *
 * 有限和 (terms 項) で近似。
 */
[[nodiscard]]
inline std::complex<double> dirichletL(const DirichletCharacter& chi, double s,
                                        int64_t terms = 10000) {
    std::complex<double> sum = 0;
    for (int64_t n = 1; n <= terms; ++n) {
        sum += chi(n) / std::pow(static_cast<double>(n), s);
    }
    return sum;
}

/**
 * @brief Gauss 和 g(χ) = Σ_{a=0}^{q-1} χ(a) · e^{2πia/q}
 */
[[nodiscard]]
inline std::complex<double> gaussSum(const DirichletCharacter& chi) {
    int64_t q = chi.modulus();
    std::complex<double> sum = 0;
    for (int64_t a = 0; a < q; ++a) {
        double angle = 2.0 * std::numbers::pi * static_cast<double>(a) / static_cast<double>(q);
        sum += chi(a) * std::complex<double>(std::cos(angle), std::sin(angle));
    }
    return sum;
}

} // namespace calx

#endif // CALX_CORE_DIRICHLET_CHARACTER_HPP
