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

/**
 * @file Diophantine.hpp
 * @brief Diophantine 方程式ソルバー (CAS-15)
 *
 * - 線形 Diophantine: ax + by = c
 * - 線形 Diophantine (多変数): a1*x1 + a2*x2 + ... = c
 * - Pell 方程式: x² - D*y² = 1
 * - 二次 Diophantine: ax² + bxy + cy² + dx + ey + f = 0
 */

#ifndef CALX_MATH_CORE_DIOPHANTINE_HPP
#define CALX_MATH_CORE_DIOPHANTINE_HPP

#include <vector>
#include <cmath>
#include <algorithm>
#include <optional>
#include <numeric>
#include <tuple>

namespace calx {

// ========================================================================
// 線形 Diophantine 方程式: ax + by = c
// ========================================================================

/// ax + by = c の整数解 (x0, y0) を求める
/// 解が存在しない場合は nullopt
/// 一般解: x = x0 + (b/g)*t, y = y0 - (a/g)*t (t は任意整数)
template<typename T = int64_t>
std::optional<std::tuple<T, T, T>> solveLinearDiophantine(T a, T b, T c) {
    if (a == 0 && b == 0) return (c == 0) ? std::make_tuple(T(0), T(0), T(0)) : std::nullopt;

    // 拡張ユークリッドの互除法: ax + by = gcd(a,b)
    auto extGcd = [](T a, T b) -> std::tuple<T, T, T> {
        if (b == 0) return {a, T(1), T(0)};
        T x0 = T(1), x1 = T(0), y0 = T(0), y1 = T(1);
        while (b != 0) {
            T q = a / b, r = a % b;
            a = b; b = r;
            T tmpX = x0 - q * x1; x0 = x1; x1 = tmpX;
            T tmpY = y0 - q * y1; y0 = y1; y1 = tmpY;
        }
        return {a, x0, y0};  // gcd, x, y
    };

    auto [g, x, y] = extGcd(std::abs(a), std::abs(b));
    if (g == 0 || c % g != 0) return std::nullopt;

    T scale = c / g;
    T x0 = x * scale * (a < 0 ? T(-1) : T(1));
    T y0 = y * scale * (b < 0 ? T(-1) : T(1));
    return std::make_tuple(x0, y0, g);
}

/// 多変数線形 Diophantine: a1*x1 + a2*x2 + ... + an*xn = c
/// 戻り値: 特殊解 x[i] (一般解は自明でない)
template<typename T = int64_t>
std::optional<std::vector<T>> solveLinearDiophantineMulti(
    const std::vector<T>& coeffs, T c)
{
    int n = static_cast<int>(coeffs.size());
    if (n == 0) return (c == 0) ? std::vector<T>{} : std::nullopt;
    if (n == 1) {
        if (coeffs[0] == 0) return (c == 0) ? std::vector<T>{T(0)} : std::nullopt;
        if (c % coeffs[0] != 0) return std::nullopt;
        return std::vector<T>{c / coeffs[0]};
    }

    // 再帰: a1*x1 + (a2*x2 + ... + an*xn) = c
    // まず gcd(a1, ..., an) | c を確認
    T g = std::abs(coeffs[0]);
    for (int i = 1; i < n; ++i) g = std::gcd(g, std::abs(coeffs[i]));
    if (g == 0) return (c == 0) ? std::vector<T>(n, T(0)) : std::nullopt;
    if (c % g != 0) return std::nullopt;

    // 2 変数に帰着して段階的に解く
    std::vector<T> result(n, T(0));
    T remaining = c;

    for (int i = 0; i < n - 1; ++i) {
        T gRest = std::abs(coeffs[i + 1]);
        for (int j = i + 2; j < n; ++j) gRest = std::gcd(gRest, std::abs(coeffs[j]));

        auto sol = solveLinearDiophantine(coeffs[i], gRest, remaining);
        if (!sol) return std::nullopt;
        auto [xi, rem_scaled, gg] = *sol;
        result[i] = xi;
        remaining = rem_scaled * gRest;
    }
    if (coeffs[n-1] != 0) {
        if (remaining % coeffs[n-1] != 0) return std::nullopt;
        result[n-1] = remaining / coeffs[n-1];
    }
    return result;
}

// ========================================================================
// Pell 方程式: x² - D*y² = 1
// ========================================================================

/// Pell 方程式 x² - D*y² = 1 の最小正の解 (x1, y1)
/// D が完全平方の場合は nullopt (解なし)
/// 連分数展開法で求解
template<typename T = int64_t>
std::optional<std::pair<T, T>> solvePell(T D, int maxIter = 10000) {
    T sqrtD = static_cast<T>(std::sqrt(static_cast<double>(D)));
    if (sqrtD * sqrtD == D) return std::nullopt;  // 完全平方

    // 連分数展開: √D = a0 + 1/(a1 + 1/(a2 + ...))
    T m = T(0), d = T(1), a0 = sqrtD, a = a0;

    // 収束列: p_{-1}=1, p_0=a0, q_{-1}=0, q_0=1
    T p_prev = T(1), p_curr = a0;
    T q_prev = T(0), q_curr = T(1);

    for (int iter = 0; iter < maxIter; ++iter) {
        m = d * a - m;
        d = (D - m * m) / d;
        if (d == 0) return std::nullopt;
        a = (a0 + m) / d;

        T p_next = a * p_curr + p_prev;
        T q_next = a * q_curr + q_prev;

        // x² - D*y² = 1 をチェック
        // オーバーフロー防止のため差分で計算
        if (p_next / q_next > 0) {  // 正の値
            T x2 = p_next * p_next;
            T Dy2 = D * q_next * q_next;
            if (x2 - Dy2 == T(1))
                return std::make_pair(p_next, q_next);
        }

        p_prev = p_curr; p_curr = p_next;
        q_prev = q_curr; q_curr = q_next;
    }
    return std::nullopt;
}

/// Pell 方程式の解列を生成 (最小解から再帰的に)
/// (x_{n+1}, y_{n+1}) = (x1*xn + D*y1*yn, x1*yn + y1*xn)
template<typename T = int64_t>
std::vector<std::pair<T, T>> pellSolutions(T D, int count, int maxIter = 10000) {
    auto base = solvePell<T>(D, maxIter);
    if (!base) return {};

    auto [x1, y1] = *base;
    std::vector<std::pair<T, T>> solutions;
    solutions.push_back({x1, y1});

    T xn = x1, yn = y1;
    for (int i = 1; i < count; ++i) {
        T xNext = x1 * xn + D * y1 * yn;
        T yNext = x1 * yn + y1 * xn;
        solutions.push_back({xNext, yNext});
        xn = xNext; yn = yNext;
    }
    return solutions;
}

// ========================================================================
// 二次 Diophantine: ax² + bxy + cy² + dx + ey + f = 0
// ========================================================================

/// 二次 Diophantine の整数解を探索 (ブルートフォース + 判別式による枝刈り)
/// 探索範囲: |x|, |y| ≤ bound
template<typename T = int64_t>
std::vector<std::pair<T, T>> solveQuadraticDiophantine(
    T a, T b, T c, T d, T e, T f, T bound = 1000)
{
    std::vector<std::pair<T, T>> solutions;

    for (T x = -bound; x <= bound; ++x) {
        // y に関する 2 次方程式: c*y² + (b*x + e)*y + (a*x² + d*x + f) = 0
        T A = c;
        T B = b * x + e;
        T C = a * x * x + d * x + f;

        if (A == 0) {
            // 1 次: By + C = 0
            if (B != 0 && C % B == 0) {
                T y = -C / B;
                solutions.push_back({x, y});
            } else if (B == 0 && C == 0) {
                solutions.push_back({x, T(0)});  // 全 y が解 (代表のみ)
            }
        } else {
            // 判別式: B² - 4AC
            T disc = B * B - T(4) * A * C;
            if (disc < 0) continue;

            T sqrtDisc = static_cast<T>(std::sqrt(static_cast<double>(disc)));
            // sqrtDisc の前後を確認 (浮動小数点誤差対策)
            for (T s = std::max(T(0), sqrtDisc - T(2)); s <= sqrtDisc + T(2); ++s) {
                if (s * s != disc) continue;
                T num1 = -B + s, num2 = -B - s;
                T denom = T(2) * A;
                if (num1 % denom == 0) solutions.push_back({x, num1 / denom});
                if (s != 0 && num2 % denom == 0) solutions.push_back({x, num2 / denom});
            }
        }
    }

    // 重複除去
    std::sort(solutions.begin(), solutions.end());
    solutions.erase(std::unique(solutions.begin(), solutions.end()), solutions.end());
    return solutions;
}

/// ピタゴラスの三つ組 (a² + b² = c²) を列挙
/// 原始的な三つ組のみ (gcd(a,b) = 1)
template<typename T = int64_t>
std::vector<std::tuple<T, T, T>> pythagoreanTriples(T maxC) {
    std::vector<std::tuple<T, T, T>> triples;
    for (T m = T(2); m * m < maxC; ++m) {
        for (T n = T(1); n < m; ++n) {
            if ((m - n) % 2 == 0) continue;  // m-n が奇数
            if (std::gcd(m, n) != 1) continue;
            T a = m * m - n * n;
            T b = T(2) * m * n;
            T c = m * m + n * n;
            if (c > maxC) break;
            if (a > b) std::swap(a, b);
            triples.push_back({a, b, c});
        }
    }
    std::sort(triples.begin(), triples.end());
    return triples;
}

} // namespace calx

#endif
