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

// BinaryQuadraticForm.hpp
//
// 二次形式 (Binary Quadratic Forms)
//
// f(x,y) = ax² + bxy + cy² の理論。
// 判別式、簡約、Gauss 合成、類数計算。
//
// 主な API:
//   BinaryQuadraticForm(a, b, c)     — 形式 ax² + bxy + cy²
//   discriminant(f)                   — 判別式 b² - 4ac
//   isReduced(f)                      — 簡約判定
//   reduce(f)                         — 簡約形
//   compose(f, g)                     — Gauss 合成
//   classNumber(D)                    — 類数 h(D)
//   representPrime(p)                 — p を表す形式

#ifndef CALX_CORE_BINARY_QUADRATIC_FORM_HPP
#define CALX_CORE_BINARY_QUADRATIC_FORM_HPP

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

namespace calx {

/**
 * @brief 二次形式 ax² + bxy + cy²
 */
class BinaryQuadraticForm {
public:
    constexpr BinaryQuadraticForm(int64_t a, int64_t b, int64_t c)
        : a_(a), b_(b), c_(c) {}

    constexpr int64_t a() const { return a_; }
    constexpr int64_t b() const { return b_; }
    constexpr int64_t c() const { return c_; }

    /// 判別式 D = b² - 4ac
    constexpr int64_t discriminant() const { return b_ * b_ - 4 * a_ * c_; }

    /// f(x, y) の評価
    constexpr int64_t evaluate(int64_t x, int64_t y) const {
        return a_ * x * x + b_ * x * y + c_ * y * y;
    }

    /// 正定値かどうか (a > 0 かつ D < 0)
    constexpr bool isPositiveDefinite() const {
        return a_ > 0 && discriminant() < 0;
    }

    /// 原始形式かどうか (gcd(a, b, c) = 1)
    bool isPrimitive() const {
        return gcd3(std::abs(a_), std::abs(b_), std::abs(c_)) == 1;
    }

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

private:
    int64_t a_, b_, c_;

    static int64_t gcd2(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 gcd3(int64_t a, int64_t b, int64_t c) {
        return gcd2(gcd2(a, b), c);
    }
};

/// 判別式 (フリー関数)
[[nodiscard]]
constexpr int64_t discriminant(const BinaryQuadraticForm& f) {
    return f.discriminant();
}

/**
 * @brief 簡約判定 (正定値の場合)
 *
 * 正定値二次形式 (D < 0) の簡約条件:
 *   |b| ≤ a ≤ c
 *   b ≥ 0 if |b| = a or a = c
 */
[[nodiscard]]
inline bool isReduced(const BinaryQuadraticForm& f) {
    if (!f.isPositiveDefinite()) return false;
    int64_t a = f.a(), b = f.b(), c = f.c();
    if (std::abs(b) > a) return false;
    if (a > c) return false;
    if (std::abs(b) == a && b < 0) return false;
    if (a == c && b < 0) return false;
    return true;
}

/**
 * @brief 簡約形 (正定値の場合)
 *
 * Gauss の簡約アルゴリズム: 隣接変換の繰り返し
 */
[[nodiscard]]
inline BinaryQuadraticForm reduce(const BinaryQuadraticForm& f) {
    int64_t a = f.a(), b = f.b(), c = f.c();

    // 正定値でない場合はそのまま返す
    if (a <= 0 || f.discriminant() >= 0) return f;

    for (int iter = 0; iter < 10000; ++iter) {
        // |b| > a なら b を簡約
        if (std::abs(b) > a) {
            // s = round(-b / (2a))
            // b' = b + 2as, c' = as² + bs + c
            int64_t s;
            if (b > 0)
                s = -(b + a) / (2 * a);
            else
                s = (-b + a) / (2 * a);

            int64_t newB = b + 2 * a * s;
            int64_t newC = a * s * s + b * s + c;
            b = newB;
            c = newC;
        }

        // a > c ならスワップ: (a,b,c) → (c,-b,a)
        if (a > c) {
            std::swap(a, c);
            b = -b;
            continue;
        }

        // |b| == a のとき b < 0 なら b を正に
        if (b == -a) { b = a; }
        // a == c のとき b < 0 なら b を正に
        if (a == c && b < 0) { b = -b; }

        break;
    }

    return {a, b, c};
}

/**
 * @brief Gauss 合成 (Dirichlet composition)
 *
 * 同じ判別式 D の二つの形式を合成する。
 * 結果も判別式 D の形式。
 */
[[nodiscard]]
inline BinaryQuadraticForm compose(const BinaryQuadraticForm& f,
                                    const BinaryQuadraticForm& g) {
    if (f.discriminant() != g.discriminant())
        throw std::invalid_argument("compose: different discriminants");

    int64_t a1 = f.a(), b1 = f.b(), c1 = f.c();
    int64_t a2 = g.a(), b2 = g.b(), c2 = g.c();

    // Shanks の合成アルゴリズム (簡易版)
    // gcd(a1, a2, (b1+b2)/2) を計算
    auto extgcd = [](int64_t a, int64_t b) -> std::tuple<int64_t, int64_t, int64_t> {
        if (a == 0) return {b, 0, 1};
        auto [g, x1, y1] = std::tuple<int64_t,int64_t,int64_t>{};
        // Iterative extended GCD
        int64_t old_r = a, r = b;
        int64_t old_s = 1, s = 0;
        int64_t old_t = 0, t = 1;
        while (r != 0) {
            int64_t q = old_r / r;
            int64_t tmp = r; r = old_r - q * r; old_r = tmp;
            tmp = s; s = old_s - q * s; old_s = tmp;
            tmp = t; t = old_t - q * t; old_t = tmp;
        }
        return {old_r, old_s, old_t};
    };

    auto [d, u, v] = extgcd(a1, a2);
    int64_t halfSum = (b1 + b2) / 2;
    auto [e, w, _] = extgcd(d, halfSum);

    int64_t a3 = (a1 / e) * (a2 / e);
    int64_t b3 = b2 + 2 * a2 * ((b1 - b2) / (2 * e) * v - (c2 * ((b1 + b2) / (2 * e))) * w);

    // b3 を mod 2*a3 の範囲に正規化
    if (a3 != 0) {
        b3 = b3 % (2 * a3);
        if (b3 > a3) b3 -= 2 * a3;
        if (b3 <= -a3) b3 += 2 * a3;
    }

    int64_t D = f.discriminant();
    int64_t c3 = (b3 * b3 - D) / (4 * a3);

    return reduce({a3, b3, c3});
}

/**
 * @brief 主形式 (principal form)
 *
 * 判別式 D の主形式: D ≡ 0 (mod 4) → (1, 0, -D/4)
 *                     D ≡ 1 (mod 4) → (1, 1, (1-D)/4)
 */
[[nodiscard]]
inline BinaryQuadraticForm principalForm(int64_t D) {
    if (D >= 0) throw std::invalid_argument("principalForm: D must be negative");
    if (D % 4 == 0)
        return {1, 0, -D / 4};
    else
        return {1, 1, (1 - D) / 4};
}

/**
 * @brief 類数 h(D)
 *
 * 判別式 D < 0 の簡約形式を全列挙して類数を求める。
 */
[[nodiscard]]
inline int64_t classNumber(int64_t D) {
    if (D >= 0) throw std::invalid_argument("classNumber: D must be negative");
    // D は判別式: D ≡ 0 or 1 (mod 4)
    if (D % 4 != 0 && D % 4 != -3)
        throw std::invalid_argument("classNumber: invalid discriminant");

    int64_t count = 0;
    // a の範囲: 1 ≤ a ≤ √(-D/3)
    int64_t aMax = static_cast<int64_t>(std::sqrt(static_cast<double>(-D) / 3.0)) + 1;

    for (int64_t a = 1; a <= aMax; ++a) {
        // |b| ≤ a, b ≡ D (mod 2)
        for (int64_t b = -a; b <= a; ++b) {
            // b² ≡ D (mod 4) → c = (b²-D)/(4a)
            int64_t num = b * b - D;
            if (num % (4 * a) != 0) continue;
            int64_t c = num / (4 * a);

            if (c < a) continue;  // a ≤ c

            BinaryQuadraticForm form(a, b, c);
            if (!isReduced(form)) continue;

            ++count;
        }
    }
    return count;
}

/**
 * @brief 簡約形式の全列挙
 */
[[nodiscard]]
inline std::vector<BinaryQuadraticForm> reducedForms(int64_t D) {
    if (D >= 0) throw std::invalid_argument("reducedForms: D must be negative");

    std::vector<BinaryQuadraticForm> forms;
    int64_t aMax = static_cast<int64_t>(std::sqrt(static_cast<double>(-D) / 3.0)) + 1;

    for (int64_t a = 1; a <= aMax; ++a) {
        for (int64_t b = -a; b <= a; ++b) {
            int64_t num = b * b - D;
            if (num % (4 * a) != 0) continue;
            int64_t c = num / (4 * a);
            if (c < a) continue;

            BinaryQuadraticForm form(a, b, c);
            if (isReduced(form)) {
                forms.push_back(form);
            }
        }
    }
    return forms;
}

} // namespace calx

#endif // CALX_CORE_BINARY_QUADRATIC_FORM_HPP
