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

// GaussianInteger.hpp
//
// ガウス整数 Z[i] = {a + bi : a, b ∈ Z}
//
// 主な API:
//   GaussianInt(a, b)      — a + bi
//   operator+, -, *, /     — 四則演算 (除算は最近接商)
//   gcd(a, b)              — ユークリッド互除法
//   norm(z)                — ノルム a² + b²
//   conj(z)                — 共役 a - bi
//   isUnit(z)              — 単元判定 (±1, ±i)
//   isPrime(z)             — ガウス素数判定
//   factorize(z)           — ガウス素因数分解

#ifndef CALX_CORE_GAUSSIAN_INTEGER_HPP
#define CALX_CORE_GAUSSIAN_INTEGER_HPP

#include <cstdint>
#include <cmath>
#include <stdexcept>
#include <vector>
#include <algorithm>
#include <utility>

namespace calx {

/**
 * @brief ガウス整数 a + bi
 */
class GaussianInt {
public:
    constexpr GaussianInt() : a_(0), b_(0) {}
    constexpr GaussianInt(int64_t a, int64_t b = 0) : a_(a), b_(b) {}

    constexpr int64_t real() const { return a_; }
    constexpr int64_t imag() const { return b_; }

    /// ノルム N(z) = a² + b²
    constexpr int64_t norm() const { return a_ * a_ + b_ * b_; }

    /// 共役
    constexpr GaussianInt conj() const { return {a_, -b_}; }

    /// 単元判定 (±1, ±i)
    constexpr bool isUnit() const { return norm() == 1; }

    constexpr bool isZero() const { return a_ == 0 && b_ == 0; }

    // 四則演算
    constexpr GaussianInt operator+(const GaussianInt& o) const {
        return {a_ + o.a_, b_ + o.b_};
    }
    constexpr GaussianInt operator-(const GaussianInt& o) const {
        return {a_ - o.a_, b_ - o.b_};
    }
    constexpr GaussianInt operator-() const { return {-a_, -b_}; }

    constexpr GaussianInt operator*(const GaussianInt& o) const {
        return {a_ * o.a_ - b_ * o.b_, a_ * o.b_ + b_ * o.a_};
    }

    /// 最近接商による除算 (z / w = round((z * conj(w)) / N(w)))
    GaussianInt operator/(const GaussianInt& w) const {
        if (w.isZero()) throw std::invalid_argument("GaussianInt: division by zero");
        // z * conj(w) = (a+bi)(c-di) = (ac+bd) + (bc-ad)i
        int64_t c = w.a_, d = w.b_;
        int64_t n = w.norm();
        // 最近接丸め除算
        int64_t re = roundDiv(a_ * c + b_ * d, n);
        int64_t im = roundDiv(b_ * c - a_ * d, n);
        return {re, im};
    }

    /// 余り
    GaussianInt operator%(const GaussianInt& w) const {
        auto q = *this / w;
        return *this - q * w;
    }

    constexpr bool operator==(const GaussianInt& o) const {
        return a_ == o.a_ && b_ == o.b_;
    }
    constexpr bool operator!=(const GaussianInt& o) const {
        return !(*this == o);
    }

private:
    int64_t a_, b_;

    /// 最近接丸め除算 (銀行家丸めではなく四捨五入)
    static int64_t roundDiv(int64_t num, int64_t den) {
        // floor((num + den/2) / den) に相当
        if (den < 0) { num = -num; den = -den; }
        if (num >= 0)
            return (2 * num + den) / (2 * den);
        else
            return -((-2 * num + den) / (2 * den));
    }
};

/// ノルム (フリー関数)
[[nodiscard]]
constexpr int64_t norm(const GaussianInt& z) { return z.norm(); }

/// 共役 (フリー関数)
[[nodiscard]]
constexpr GaussianInt conj(const GaussianInt& z) { return z.conj(); }

/**
 * @brief ガウス整数の GCD (ユークリッド互除法)
 *
 * 結果は associate (単元倍の自由度あり) まで一意。
 */
[[nodiscard]]
inline GaussianInt gcd(GaussianInt a, GaussianInt b) {
    while (!b.isZero()) {
        auto r = a % b;
        a = b;
        b = r;
    }
    return a;
}

/**
 * @brief ガウス素数判定
 *
 * z がガウス素数 ⟺
 *   - z = a + 0i: |a| が通常の素数で a ≡ 3 (mod 4)
 *   - z = 0 + bi: |b| が通常の素数で b ≡ 3 (mod 4)
 *   - N(z) が通常の素数
 */
[[nodiscard]]
inline bool isGaussianPrime(const GaussianInt& z) {
    if (z.isZero()) return false;
    if (z.isUnit()) return false;

    int64_t n = z.norm();

    // N(z) が通常の素数なら z はガウス素数
    // 通常素数判定
    auto isPrime = [](int64_t n) -> bool {
        if (n < 2) return false;
        if (n < 4) return true;
        if (n % 2 == 0 || n % 3 == 0) return false;
        for (int64_t i = 5; i * i <= n; i += 6)
            if (n % i == 0 || n % (i + 2) == 0) return false;
        return true;
    };

    // ケース1: 純実数または純虚数
    if (z.imag() == 0) {
        int64_t a = std::abs(z.real());
        return isPrime(a) && (a % 4 == 3);
    }
    if (z.real() == 0) {
        int64_t b = std::abs(z.imag());
        return isPrime(b) && (b % 4 == 3);
    }

    // ケース2: N(z) が素数
    return isPrime(n);
}

/**
 * @brief ガウス整数の素因数分解
 *
 * 結果は (ガウス素数, 指数) のペアのリスト。
 * 単元因子は含まない (associate まで一意)。
 */
[[nodiscard]]
inline std::vector<std::pair<GaussianInt, int>> factorize(GaussianInt z) {
    if (z.isZero()) return {};
    if (z.isUnit()) return {};

    std::vector<std::pair<GaussianInt, int>> factors;

    auto isPrime = [](int64_t n) -> bool {
        if (n < 2) return false;
        if (n < 4) return true;
        if (n % 2 == 0 || n % 3 == 0) return false;
        for (int64_t i = 5; i * i <= n; i += 6)
            if (n % i == 0 || n % (i + 2) == 0) return false;
        return true;
    };

    // ノルムを素因数分解して、各素因数に対してガウス素因数を見つける
    int64_t n = z.norm();

    // 2 = -i(1+i)² → 1+i で割る
    {
        GaussianInt pi(1, 1);
        int exp = 0;
        while (!z.isZero()) {
            auto r = z % pi;
            if (r.isZero()) {
                z = z / pi;
                ++exp;
            } else break;
        }
        if (exp > 0) factors.push_back({pi, exp});
    }

    // 奇素数 p
    for (int64_t p = 3; p * p <= n; p += 2) {
        if (!isPrime(p)) continue;
        if (n % p != 0) continue;

        if (p % 4 == 3) {
            // p はガウスでも素: p | z ⟺ p | N(z)
            GaussianInt gp(p, 0);
            int exp = 0;
            while (!z.isZero()) {
                auto r = z % gp;
                if (r.isZero()) {
                    z = z / gp;
                    ++exp;
                } else break;
            }
            if (exp > 0) factors.push_back({gp, exp});
        } else {
            // p ≡ 1 (mod 4): p = π · conj(π) で分解
            // a² ≡ -1 (mod p) を解いて π = gcd(a + i, p)
            int64_t a = 0;
            for (int64_t x = 2; x < p; ++x) {
                // x^((p-1)/2) mod p
                int64_t result = 1, base = x % p, exp = (p - 1) / 2;
                int64_t tmp = base;
                int64_t e = exp;
                result = 1;
                while (e > 0) {
                    if (e & 1) result = (result * tmp) % p;
                    tmp = (tmp * tmp) % p;
                    e >>= 1;
                }
                if (result == p - 1) {
                    // x は非剰余 → x^((p-1)/4) mod p が √-1
                    int64_t r2 = 1;
                    tmp = x % p;
                    e = (p - 1) / 4;
                    while (e > 0) {
                        if (e & 1) r2 = (r2 * tmp) % p;
                        tmp = (tmp * tmp) % p;
                        e >>= 1;
                    }
                    a = r2;
                    break;
                }
            }

            GaussianInt pi = gcd(GaussianInt(a, 1), GaussianInt(p, 0));
            GaussianInt piConj = pi.conj();

            // pi で割れるだけ割る
            int exp1 = 0;
            while (!z.isZero()) {
                auto r = z % pi;
                if (r.isZero()) { z = z / pi; ++exp1; }
                else break;
            }
            if (exp1 > 0) factors.push_back({pi, exp1});

            // conj(pi) でも
            int exp2 = 0;
            while (!z.isZero()) {
                auto r = z % piConj;
                if (r.isZero()) { z = z / piConj; ++exp2; }
                else break;
            }
            if (exp2 > 0) factors.push_back({piConj, exp2});
        }
    }

    // 残りが 1 でなければそれ自体が素因数
    if (!z.isUnit() && !z.isZero()) {
        factors.push_back({z, 1});
    }

    return factors;
}

} // namespace calx

#endif // CALX_CORE_GAUSSIAN_INTEGER_HPP
