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

// ConvexOptimization_impl.hpp — 凸最適化ソルバーの実装
// proxL1, proxNonneg, fista, lassoFISTA, admmLasso

#ifndef SANGI_MATH_OPTIMIZATION_CONVEX_IMPL_HPP
#define SANGI_MATH_OPTIMIZATION_CONVEX_IMPL_HPP

#include <math/optimization/ConvexOptimization.hpp>

namespace sangi {

    // ================================================================
    // 近接作用素 (Proximal Operators)
    // ================================================================

    template<typename T>
    Vector<T> proxL1(const Vector<T>& v, T lambda) {
        size_t n = v.size();
        Vector<T> result(n);
        for (size_t i = 0; i < n; ++i) {
            T vi = v[i];
            if (vi > lambda) result[i] = vi - lambda;
            else if (vi < -lambda) result[i] = vi + lambda;
            else result[i] = T{0};
        }
        return result;
    }

    template<typename T>
    Vector<T> proxNonneg(const Vector<T>& v) {
        size_t n = v.size();
        Vector<T> result(n);
        for (size_t i = 0; i < n; ++i)
            result[i] = std::max(v[i], T{0});
        return result;
    }

    // ================================================================
    // FISTA
    // ================================================================

    template<typename T>
    FISTAResult<T> fista(
        std::function<Vector<T>(const Vector<T>&)> gradF,
        std::function<Vector<T>(const Vector<T>&, T)> proxG,
        const Vector<T>& x0,
        T L,
        size_t maxIter, T tol)
    {
        size_t n = x0.size();
        Vector<T> x = x0;
        Vector<T> y = x0;
        T t = T{1};
        T stepSize = T{1} / L;

        for (size_t iter = 0; iter < maxIter; ++iter) {
            Vector<T> xPrev = x;

            // 勾配ステップ + 近接ステップ
            auto grad = gradF(y);
            Vector<T> z(n);
            for (size_t i = 0; i < n; ++i)
                z[i] = y[i] - stepSize * grad[i];
            x = proxG(z, stepSize);

            // FISTA の加速 (Nesterov momentum)
            T tNew = (T{1} + std::sqrt(T{1} + T{4} * t * t)) / T{2};
            T beta = (t - T{1}) / tNew;
            for (size_t i = 0; i < n; ++i)
                y[i] = x[i] + beta * (x[i] - xPrev[i]);
            t = tNew;

            // 収束判定
            T diff = T{0};
            for (size_t i = 0; i < n; ++i) {
                T d = x[i] - xPrev[i];
                diff += d * d;
            }
            if (std::sqrt(diff) < tol) {
                return { std::move(x), T{0}, iter + 1 };
            }
        }
        return { std::move(x), T{0}, maxIter };
    }

    template<typename T>
    FISTAResult<T> lassoFISTA(const Matrix<T>& A, const Vector<T>& b,
                               T lambda, size_t maxIter, T tol)
    {
        size_t m = A.rows(), n = A.cols();

        // A^T A の最大固有値の近似 (パワーメソッド)
        Vector<T> v(n, T{1});
        T L = T{1};
        for (int iter = 0; iter < 30; ++iter) {
            // w = A^T (A v)
            Vector<T> Av(m, T{0});
            for (size_t i = 0; i < m; ++i)
                for (size_t j = 0; j < n; ++j)
                    Av[i] += A(i, j) * v[j];
            Vector<T> w(n, T{0});
            for (size_t j = 0; j < n; ++j)
                for (size_t i = 0; i < m; ++i)
                    w[j] += A(i, j) * Av[i];
            T norm = T{0};
            for (auto x : w) norm += x * x;
            L = std::sqrt(norm);
            if (L > T{1e-15})
                for (size_t j = 0; j < n; ++j) v[j] = w[j] / L;
        }

        auto gradF = [&](const Vector<T>& x) -> Vector<T> {
            // ∇f = A^T(Ax - b)
            Vector<T> Ax(m, T{0});
            for (size_t i = 0; i < m; ++i)
                for (size_t j = 0; j < n; ++j)
                    Ax[i] += A(i, j) * x[j];
            for (size_t i = 0; i < m; ++i) Ax[i] -= b[i];
            Vector<T> grad(n, T{0});
            for (size_t j = 0; j < n; ++j)
                for (size_t i = 0; i < m; ++i)
                    grad[j] += A(i, j) * Ax[i];
            return grad;
        };

        auto proxG = [&](const Vector<T>& v, T step) -> Vector<T> {
            return proxL1(v, lambda * step);
        };

        Vector<T> x0(n, T{0});
        return fista(
            std::function<Vector<T>(const Vector<T>&)>(gradF),
            std::function<Vector<T>(const Vector<T>&, T)>(proxG),
            x0, L, maxIter, tol);
    }

    // ================================================================
    // ADMM
    // ================================================================

    template<typename T>
    ADMMResult<T> admmLasso(const Matrix<T>& A, const Vector<T>& b,
                             T lambda, T rho,
                             size_t maxIter, T tol)
    {
        size_t m = A.rows(), n = A.cols();

        // (A^T A + ρ I)^{-1} を事前計算 (Cholesky で解くべきだが簡略版)
        // 小規模では直接逆行列
        Matrix<T> AtA(n, n, T{0});
        Vector<T> Atb(n, T{0});
        for (size_t j = 0; j < n; ++j) {
            for (size_t i = 0; i < m; ++i) {
                Atb[j] += A(i, j) * b[i];
                for (size_t k = 0; k < n; ++k)
                    AtA(j, k) += A(i, j) * A(i, k);
            }
            AtA(j, j) += rho;
        }

        Vector<T> x(n, T{0}), z(n, T{0}), u(n, T{0});

        for (size_t iter = 0; iter < maxIter; ++iter) {
            // x ← (A^T A + ρI)^{-1} (A^T b + ρ(z - u))
            Vector<T> rhs(n);
            for (size_t j = 0; j < n; ++j)
                rhs[j] = Atb[j] + rho * (z[j] - u[j]);

            // 簡易ソルバー (ガウスの消去法)
            Matrix<T> Aug(n, n + 1);
            for (size_t i = 0; i < n; ++i) {
                for (size_t j = 0; j < n; ++j)
                    Aug(i, j) = AtA(i, j);
                Aug(i, n) = rhs[i];
            }
            for (size_t k = 0; k < n; ++k) {
                T pivot = Aug(k, k);
                if (std::abs(pivot) < T{1e-15}) continue;
                for (size_t j = k; j <= n; ++j) Aug(k, j) /= pivot;
                for (size_t i = 0; i < n; ++i) {
                    if (i == k) continue;
                    T factor = Aug(i, k);
                    for (size_t j = k; j <= n; ++j)
                        Aug(i, j) -= factor * Aug(k, j);
                }
            }
            for (size_t i = 0; i < n; ++i) x[i] = Aug(i, n);

            // z ← prox_{λ/ρ ||·||₁}(x + u)
            Vector<T> xpu(n);
            for (size_t i = 0; i < n; ++i) xpu[i] = x[i] + u[i];
            z = proxL1(xpu, lambda / rho);

            // u ← u + x - z
            T primalRes = T{0};
            for (size_t i = 0; i < n; ++i) {
                u[i] += x[i] - z[i];
                primalRes += (x[i] - z[i]) * (x[i] - z[i]);
            }
            if (std::sqrt(primalRes) < tol)
                return { std::move(x), iter + 1 };
        }
        return { std::move(x), maxIter };
    }

} // namespace sangi

#endif // SANGI_MATH_OPTIMIZATION_CONVEX_IMPL_HPP