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

// root_finding_nd.hpp
#ifndef SANGI_ROOT_FINDING_ND_HPP
#define SANGI_ROOT_FINDING_ND_HPP

#include "root_finding_base.hpp"
#include <math/core/vector.hpp>
#include <math/core/matrix.hpp>

namespace sangi {

/**
 * @brief 多次元問題の求根結果を表す構造体
 */
template<typename V, typename T = typename V::value_type>
struct SystemRootFindingResult {
    std::optional<V> root;  // 根の近似値（失敗した場合はnullopt）
    bool converged;         // 収束したかどうか
    size_t iterations;      // 実行した反復回数
    T error_estimate;       // 誤差の推定値
    
    // 成功結果を作成するヘルパー
    static SystemRootFindingResult success(const V& root_value, size_t iter_count, T err = T(0)) {
        return {root_value, true, iter_count, err};
    }
    
    // 失敗結果を作成するヘルパー
    static SystemRootFindingResult failure(size_t iter_count, T err = T(0)) {
        return {std::nullopt, false, iter_count, err};
    }
    
    // オプショナル成功結果（条件付き成功）
    static SystemRootFindingResult partial_success(const V& root_value, bool conv, size_t iter_count, T err = T(0)) {
        return {root_value, conv, iter_count, err};
    }
};

/**
 * @brief 非線形方程式系をニュートン法で解く
 * @tparam T 順序構造を持つ体型（実数など）
 * @tparam V ベクトル型
 * @tparam M 行列型
 * @param F 方程式系のベクトル値関数
 * @param J ヤコビアン行列を計算する関数
 * @param x0 初期値ベクトル
 * @param criteria 収束判定基準
 * @return 根のベクトルと収束状態を含む結果オブジェクト
 */
template<concepts::OrderedField T, typename V, typename M>
    requires concepts::VectorOf<V, T> && concepts::MatrixOf<M, T> && requires(M m, V v) {
        { m.solve(v) } -> std::convertible_to<std::optional<V>>;
    }
SystemRootFindingResult<V, T> newton_raphson_system(
    const std::function<V(const V&)>& F,
    const std::function<M(const V&)>& J,
    const V& x0,
    const ConvergenceCriteria<T>& criteria = ConvergenceCriteria<T>()) {

    V x = x0;
    V x_prev = x0;
    size_t iterations = 0;

    for (iterations = 0; iterations < criteria.max_iterations; ++iterations) {
        // 現在点での関数値とヤコビアンを計算
        V f = F(x);

        // 関数値ノルムによる収束判定
        T f_norm = norm2(f);
        if (f_norm < criteria.abs_ftol) {
            return SystemRootFindingResult<V, T>::success(x, iterations, f_norm);
        }

        M jacobian = J(x);
        auto dx_opt = jacobian.solve(-f);

        // 特異行列対策：対角成分に小さな値を追加して正則化
        if (!dx_opt) {
            T reg_factor = criteria.abs_xtol * 10;
            M reg_jacobian = jacobian;
            for (size_t j = 0; j < reg_jacobian.rows(); ++j) {
                reg_jacobian(j, j) += reg_factor;
            }
            dx_opt = reg_jacobian.solve(-f);
            if (!dx_opt) {
                return SystemRootFindingResult<V, T>::failure(iterations, f_norm);
            }
        }

        V dx = *dx_opt;

        // 発散防止策：ステップサイズの制限
        T step_size = norm2(dx);
        const T max_step = T(10);  // 最大ステップサイズ
        if (step_size > max_step) {
            dx = dx * (max_step / step_size);
        }

        x_prev = x;
        V x_new = x + dx;
        x = x_new;

        // ベクトル値による収束判定
        if (criteria.is_vector_converged(f, x, x_prev)) {
            return SystemRootFindingResult<V, T>::success(x, iterations + 1, f_norm);
        }
    }

    // 最大反復回数に達したが、最後の近似値を返す
    V f_final = F(x);
    T error = norm2(f_final);
    bool converged = error < criteria.abs_ftol * 10; // 緩和された収束判定
    
    return SystemRootFindingResult<V, T>::partial_success(x, converged, iterations, error);
}

/**
 * @brief ブロイデン法による非線形方程式系の求根
 * @tparam T 順序構造を持つ体型（実数など）
 * @tparam V ベクトル型
 * @tparam M 行列型
 * @param F 方程式系のベクトル値関数
 * @param x0 初期値ベクトル
 * @param J0 初期ヤコビアン行列（省略可）
 * @param criteria 収束判定基準
 * @return 根のベクトルと収束状態を含む結果オブジェクト
 */
template<concepts::OrderedField T, typename V, typename M>
    requires concepts::VectorOf<V, T> && concepts::MatrixOf<M, T> && requires(M m, V v) {
        { m.solve(v) } -> std::convertible_to<std::optional<V>>;
    }
SystemRootFindingResult<V, T> broyden_method(
    const std::function<V(const V&)>& F,
    const V& x0,
    std::optional<M> J0 = std::nullopt,
    const ConvergenceCriteria<T>& criteria = ConvergenceCriteria<T>()) {

    const size_t n = x0.size();
    V x = x0;
    V f = F(x);
    T f_norm = norm2(f);
    
    // 初期ヤコビアンの設定
    M B;  // 近似ヤコビアン
    if (J0) {
        B = *J0;
    } else {
        // 初期ヤコビアンを単位行列として設定
        B = identity<T>(n);
    }
    
    size_t iterations = 0;
    
    for (iterations = 0; iterations < criteria.max_iterations; ++iterations) {
        // 収束判定
        if (f_norm < criteria.abs_ftol) {
            return SystemRootFindingResult<V, T>::success(x, iterations, f_norm);
        }
        
        // ステップを計算
        auto dx_opt = B.solve(-f);
        if (!dx_opt) {
            // 特異行列対策
            T reg_factor = criteria.abs_xtol * 10;
            M reg_B = B;
            for (size_t j = 0; j < n; ++j) {
                reg_B(j, j) += reg_factor;
            }
            dx_opt = reg_B.solve(-f);
            if (!dx_opt) {
                return SystemRootFindingResult<V, T>::failure(iterations, f_norm);
            }
        }
        
        V dx = *dx_opt;
        
        // ステップサイズの制限
        T step_size = norm2(dx);
        const T max_step = T(10);
        if (step_size > max_step) {
            dx = dx * (max_step / step_size);
        }
        
        V x_new = x + dx;
        V f_new = F(x_new);
        T f_new_norm = norm2(f_new);
        
        // ライン探索的な調整
        T alpha = T(1);
        if (f_new_norm > f_norm) {
            // バックトラッキング
            const T min_alpha = T(0.1);
            const T reduce_factor = T(0.5);

            while (f_new_norm > f_norm && alpha > min_alpha) {
                alpha *= reduce_factor;
                x_new = x + alpha * dx;
                f_new = F(x_new);
                f_new_norm = norm2(f_new);
            }

            if (f_new_norm > f_norm) {
                // 改善が見られない場合
                return SystemRootFindingResult<V, T>::partial_success(x, false, iterations, f_norm);
            }
        }

        // ブロイデン更新
        V df = f_new - f;
        V s = alpha * dx;  // 実際のステップ
        
        // ブロイデン更新の分母チェック
        T s_dot_df = dot(s, df);
        if (std::abs(s_dot_df) > std::numeric_limits<T>::epsilon()) {
            V Bs = B * s;
            V temp = (df - Bs) / s_dot_df;
            
            // ランク1更新
            for (size_t i = 0; i < n; ++i) {
                for (size_t j = 0; j < n; ++j) {
                    B(i, j) += temp[i] * s[j];
                }
            }
        }
        
        // 更新
        x = x_new;
        f = f_new;
        f_norm = f_new_norm;
        
        // 収束判定 - 改めてチェック
        if (f_norm < criteria.abs_ftol) {
            return SystemRootFindingResult<V, T>::success(x, iterations + 1, f_norm);
        }
    }
    
    // 最大反復回数に達した場合
    bool converged = f_norm < criteria.abs_ftol * 10;
    return SystemRootFindingResult<V, T>::partial_success(x, converged, iterations, f_norm);
}

/**
 * @brief 非線形最小二乗法（レーベンバーグ・マーカート法）
 * @tparam T 順序構造を持つ体型（実数など）
 * @tparam V ベクトル型
 * @tparam M 行列型
 * @param F 残差ベクトルを返す関数
 * @param J ヤコビアン行列を計算する関数
 * @param x0 初期値ベクトル
 * @param criteria 収束判定基準
 * @return 最適解のベクトルと収束状態を含む結果オブジェクト
 */
template<concepts::OrderedField T, typename V, typename M>
    requires concepts::VectorOf<V, T> && concepts::MatrixOf<M, T> && requires(M m, V v) {
        { m.solve(v) } -> std::convertible_to<std::optional<V>>;
    }
SystemRootFindingResult<V, T> levenberg_marquardt(
    const std::function<V(const V&)>& F,
    const std::function<M(const V&)>& J,
    const V& x0,
    const ConvergenceCriteria<T>& criteria = ConvergenceCriteria<T>()) {

    V x = x0;
    V f = F(x);
    T f_norm = norm2(f);
    T f_squared = dot(f, f);  // 残差二乗和
    
    // LM法のダンピングパラメータ
    T lambda = T(0.01);
    const T lambda_down = T(0.1);  // λ減少因子
    const T lambda_up = T(10);     // λ増加因子
    
    size_t iterations = 0;
    size_t consecutive_failures = 0;
    const size_t max_failures = 5;  // 連続失敗の最大回数
    
    for (iterations = 0; iterations < criteria.max_iterations; ++iterations) {
        // 収束判定
        if (f_norm < criteria.abs_ftol) {
            return SystemRootFindingResult<V, T>::success(x, iterations, f_norm);
        }
        
        // ヤコビアン計算
        M jac = J(x);
        M JtJ = transpose(jac) * jac;
        V JtF = transpose(jac) * f;
        
        // LM方程式の構築
        M A = JtJ;
        for (size_t i = 0; i < A.rows(); ++i) {
            A(i, i) += lambda * (A(i, i) + T(1));  // 対角要素を強化
        }
        
        // ステップの計算
        auto dx_opt = A.solve(-JtF);
        if (!dx_opt) {
            consecutive_failures++;
            if (consecutive_failures >= max_failures) {
                return SystemRootFindingResult<V, T>::partial_success(x, false, iterations, f_norm);
            }
            
            // λを増加させて再試行
            lambda *= lambda_up;
            continue;
        }
        
        V dx = *dx_opt;
        
        // 試行的なステップ
        V x_new = x + dx;
        V f_new = F(x_new);
        T f_new_squared = dot(f_new, f_new);
        
        // ステップの評価
        if (f_new_squared < f_squared) {
            // 成功：ステップを受け入れ、λを減少
            x = x_new;
            f = f_new;
            f_norm = norm2(f);
            f_squared = f_new_squared;
            lambda *= lambda_down;
            consecutive_failures = 0;
        } else {
            // 失敗：ステップを拒否、λを増加
            lambda *= lambda_up;
            consecutive_failures++;
            
            if (consecutive_failures >= max_failures) {
                return SystemRootFindingResult<V, T>::partial_success(x, false, iterations, f_norm);
            }
        }
    }
    
    // 最大反復回数に達した場合
    bool converged = f_norm < criteria.abs_ftol * 10;
    return SystemRootFindingResult<V, T>::partial_success(x, converged, iterations, f_norm);
}

/**
 * @brief Powell ハイブリッド法 (MINPACK hybrd 相当)
 *
 * 有限差分ヤコビアン + dogleg trust region + Broyden rank-1 更新を組み合わせた
 * 頑健な非線形方程式系ソルバー。ヤコビアンの解析的導関数が不要。
 *
 * Newton 法が適用可能な場合は Newton ステップを使用し、
 * ヤコビアンが不良条件や特異の場合は scaled gradient (Cauchy) ステップに退化する。
 * Trust region により大域的収束性を確保。
 *
 * @tparam T  順序体型
 * @tparam V  ベクトル型 (operator[], size(), 算術演算)
 * @tparam M  行列型 (operator(i,j), rows(), cols())
 * @param F        方程式系 F(x) = 0
 * @param x0       初期値ベクトル
 * @param criteria 収束判定基準
 * @return 根のベクトルと収束状態
 */
template<concepts::OrderedField T, typename V, typename M>
    requires concepts::VectorOf<V, T> && concepts::MatrixOf<M, T>
SystemRootFindingResult<V, T> powell_hybrid(
    const std::function<V(const V&)>& F,
    const V& x0,
    const ConvergenceCriteria<T>& criteria = ConvergenceCriteria<T>())
{
    const size_t n = x0.size();
    V x = x0;
    V f = F(x);
    T f_norm = norm2(f);

    // 内部ガウス消去ソルバー: A*x = b (部分ピボット)
    auto solve_linear = [&](M A, V b) -> std::optional<V> {
        for (size_t k = 0; k < n; ++k) {
            // ピボット選択
            size_t max_row = k;
            T max_val = std::abs(A(k, k));
            for (size_t i = k + 1; i < n; ++i) {
                T v = std::abs(A(i, k));
                if (v > max_val) { max_val = v; max_row = i; }
            }
            if (max_val < std::numeric_limits<T>::epsilon() * T(100))
                return std::nullopt;
            if (max_row != k) {
                for (size_t j = k; j < n; ++j) std::swap(A(max_row, j), A(k, j));
                std::swap(b[max_row], b[k]);
            }
            // 前進消去
            for (size_t i = k + 1; i < n; ++i) {
                T factor = A(i, k) / A(k, k);
                for (size_t j = k + 1; j < n; ++j)
                    A(i, j) -= factor * A(k, j);
                b[i] -= factor * b[k];
            }
        }
        // 後退代入
        V result = b;
        for (int k = static_cast<int>(n) - 1; k >= 0; --k) {
            for (size_t j = static_cast<size_t>(k) + 1; j < n; ++j)
                result[k] -= A(k, j) * result[j];
            result[k] /= A(k, k);
        }
        return result;
    };

    // 有限差分ヤコビアン計算
    auto compute_jacobian = [&](const V& xc, const V& fc) -> M {
        M J = identity<T>(n);
        T eps_fd = std::sqrt(std::numeric_limits<T>::epsilon());
        for (size_t j = 0; j < n; ++j) {
            V xp = xc;
            T h = eps_fd * std::max(std::abs(xc[j]), T(1));
            xp[j] += h;
            V fp = F(xp);
            for (size_t i = 0; i < n; ++i) {
                J(i, j) = (fp[i] - fc[i]) / h;
            }
        }
        return J;
    };

    // ヤコビアン計算
    M J = compute_jacobian(x, f);

    // Trust region 半径
    T delta = T(100) * norm2(x);
    if (delta < T(1)) delta = T(100);
    const T delta_max = T(1e8);

    // 対角スケーリングベクトル
    V diag = x;
    for (size_t j = 0; j < n; ++j) {
        T col_norm = T(0);
        for (size_t i = 0; i < n; ++i)
            col_norm += J(i, j) * J(i, j);
        diag[j] = std::max(std::sqrt(col_norm), T(1));
    }

    size_t iterations = 0;
    size_t jacobian_age = 0;
    const size_t jacobian_refresh = n * 2;

    for (iterations = 0; iterations < criteria.max_iterations; ++iterations) {
        // 収束判定
        if (f_norm < criteria.abs_ftol) {
            return SystemRootFindingResult<V, T>::success(x, iterations, f_norm);
        }

        // Newton ステップ: J * p_newton = -f
        V neg_f = f * T(-1);
        auto p_newton_opt = solve_linear(J, neg_f);
        V p_newton;
        bool newton_ok = false;
        T p_newton_norm = T(0);

        if (p_newton_opt) {
            p_newton = *p_newton_opt;
            p_newton_norm = norm2(p_newton);
            newton_ok = std::isfinite(p_newton_norm) && p_newton_norm < T(1e15);
        }

        // Cauchy ステップ (scaled gradient): p_cauchy = -α * J^T * f
        V jtf = x;
        for (size_t j = 0; j < n; ++j) {
            T s = T(0);
            for (size_t i = 0; i < n; ++i) s += J(i, j) * f[i];
            jtf[j] = s;
        }
        T jtf_norm = norm2(jtf);

        V jjtf = f;
        for (size_t i = 0; i < n; ++i) {
            T s = T(0);
            for (size_t j = 0; j < n; ++j) s += J(i, j) * jtf[j];
            jjtf[i] = s;
        }
        T jjtf_norm = norm2(jjtf);

        T alpha = (jjtf_norm > T(0))
            ? (jtf_norm * jtf_norm) / (jjtf_norm * jjtf_norm)
            : T(1);
        V p_cauchy = jtf * (-alpha);
        T p_cauchy_norm = norm2(p_cauchy);

        // Dogleg ステップ選択
        V step = p_cauchy;
        if (newton_ok && p_newton_norm <= delta) {
            // Newton ステップが trust region 内 → そのまま使用
            step = p_newton;
        } else if (p_cauchy_norm >= delta) {
            // Cauchy ステップですら trust region 外 → スケーリング
            step = p_cauchy * (delta / p_cauchy_norm);
        } else if (newton_ok) {
            // Dogleg 補間: Cauchy → Newton の直線上で delta と交わる点
            V diff = p_newton - p_cauchy;
            T diff_norm_sq = dot(diff, diff);
            T pc_dot_diff = dot(p_cauchy, diff);
            T pc_norm_sq = p_cauchy_norm * p_cauchy_norm;

            // ||p_cauchy + tau * diff||² = delta² を解く
            T discriminant = pc_dot_diff * pc_dot_diff
                - diff_norm_sq * (pc_norm_sq - delta * delta);
            if (discriminant >= T(0)) {
                T tau = (-pc_dot_diff + std::sqrt(discriminant)) / diff_norm_sq;
                tau = std::max(T(0), std::min(T(1), tau));
                step = p_cauchy + diff * tau;
            } else {
                step = p_cauchy * (delta / p_cauchy_norm);
            }
        }

        // 試行ステップ
        V x_new = x + step;
        V f_new = F(x_new);
        T f_new_norm = norm2(f_new);

        // 実際の減少 vs 予測減少
        T actual_reduction = f_norm * f_norm - f_new_norm * f_new_norm;
        // 予測: ||f||² - ||f + J*step||²
        V jstep = f;
        for (size_t i = 0; i < n; ++i) {
            T s = T(0);
            for (size_t j = 0; j < n; ++j) s += J(i, j) * step[j];
            jstep[i] = f[i] + s;
        }
        T predicted_norm_sq = dot(jstep, jstep);
        T predicted_reduction = f_norm * f_norm - predicted_norm_sq;

        T ratio = (std::abs(predicted_reduction) > std::numeric_limits<T>::epsilon() * f_norm * f_norm)
            ? actual_reduction / predicted_reduction
            : T(0);

        // Trust region 更新
        if (ratio < T(0.1)) {
            delta *= T(0.5);
            if (delta < criteria.abs_xtol) {
                // Trust region が極小 → 収束とみなす
                return SystemRootFindingResult<V, T>::partial_success(
                    x, f_norm < criteria.abs_ftol * T(10), iterations, f_norm);
            }
        } else if (ratio > T(0.75)) {
            delta = std::min(T(2) * delta, delta_max);
        }

        // ステップ受理判定
        if (ratio > T(0.0001)) {
            // Broyden rank-1 更新: J_new = J + (df - J*s)*s^T / (s^T*s)
            V s = step;
            V df = f_new - f;
            T sts = dot(s, s);

            if (sts > std::numeric_limits<T>::epsilon()) {
                // J*s
                V js = f;
                for (size_t i = 0; i < n; ++i) {
                    T sum = T(0);
                    for (size_t j = 0; j < n; ++j) sum += J(i, j) * s[j];
                    js[i] = sum;
                }
                // (df - J*s) / (s^T*s)
                V update = (df - js) * (T(1) / sts);
                for (size_t i = 0; i < n; ++i)
                    for (size_t j = 0; j < n; ++j)
                        J(i, j) += update[i] * s[j];
            }

            x = x_new;
            f = f_new;
            f_norm = f_new_norm;
            ++jacobian_age;

            // 対角スケーリング更新
            for (size_t j = 0; j < n; ++j) {
                T col_norm = T(0);
                for (size_t i = 0; i < n; ++i)
                    col_norm += J(i, j) * J(i, j);
                diag[j] = std::max(diag[j], std::sqrt(col_norm));
            }

            // ヤコビアンリフレッシュ
            if (jacobian_age >= jacobian_refresh) {
                J = compute_jacobian(x, f);
                jacobian_age = 0;
            }

            // 収束判定
            if (f_norm < criteria.abs_ftol) {
                return SystemRootFindingResult<V, T>::success(x, iterations + 1, f_norm);
            }
        }
        // ratio が小さすぎる場合: ステップ棄却、delta は既に縮小済み
    }

    bool converged = f_norm < criteria.abs_ftol * T(10);
    return SystemRootFindingResult<V, T>::partial_success(x, converged, iterations, f_norm);
}

// 旧インターフェースとの互換性のためのラッパー関数
template<concepts::OrderedField T, typename V, typename M>
    requires concepts::VectorOf<V, T> && concepts::MatrixOf<M, T>
std::optional<V> newton_raphson_system_legacy(
    const std::function<V(const V&)>& F,
    const std::function<M(const V&)>& J,
    const V& x0,
    T tolerance = std::numeric_limits<T>::epsilon() * 100,
    size_t max_iterations = 50) {

    ConvergenceCriteria<T> criteria;
    criteria.abs_ftol = tolerance;
    criteria.abs_xtol = tolerance;
    criteria.max_iterations = max_iterations;

    auto result = newton_raphson_system<T, V, M>(F, J, x0, criteria);
    return result.root;
}

#if SANGI_HAS_MKL
// MKL実装をここに追加
#endif

} // namespace sangi

#endif // SANGI_ROOT_FINDING_ND_HPP