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

// ConvexOptimization_impl.hpp — implementation of the convex optimization solvers
// proxL1, proxNonneg, fista, lassoFISTA, admmLasso, solveQP

#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 BaseVector<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 BaseVector<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;

            // gradient step + proximal step
            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 acceleration (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;

            // convergence check
            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 BaseMatrix<T>& A, const BaseVector<T>& b,
                               T lambda, size_t maxIter, T tol)
    {
        size_t m = A.rows(), n = A.cols();

        // approximation of the largest eigenvalue of A^T A (power method)
        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 BaseMatrix<T>& A, const BaseVector<T>& b,
                             T lambda, T rho,
                             size_t maxIter, T tol)
    {
        size_t m = A.rows(), n = A.cols();

        // precompute (A^T A + ρ I)^{-1} (should solve via Cholesky, but this is a simplified version)
        // direct inverse for small problems
        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]);

            // simple solver (Gaussian elimination)
            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 };
    }

    // ================================================================
    // QP (interior-point method)
    // ================================================================

    template<typename T>
    QPResult<T> solveQP(const BaseMatrix<T>& Q, const BaseVector<T>& c,
                         const BaseMatrix<T>& A, const BaseVector<T>& b,
                         size_t maxIter, T tol)
    {
        size_t n = c.size();       // number of variables
        size_t m = b.size();       // number of constraints

        // initial feasible point
        Vector<T> x(n, T{1});
        Vector<T> s(m, T{1});      // slacks
        Vector<T> lambda(m, T{1}); // Lagrange multipliers

        for (size_t iter = 0; iter < maxIter; ++iter) {
            // residual computation
            // r_dual = Qx + c - A^T λ
            Vector<T> rDual(n, T{0});
            for (size_t j = 0; j < n; ++j) {
                rDual[j] = c[j];
                for (size_t k = 0; k < n; ++k) rDual[j] += Q(j, k) * x[k];
                for (size_t i = 0; i < m; ++i) rDual[j] -= A(i, j) * lambda[i];
            }

            // r_primal = Ax + s - b
            Vector<T> rPrimal(m, T{0});
            for (size_t i = 0; i < m; ++i) {
                rPrimal[i] = -b[i] + s[i];
                for (size_t j = 0; j < n; ++j) rPrimal[i] += A(i, j) * x[j];
            }

            // complementarity μ = s^T λ / m
            T mu = T{0};
            for (size_t i = 0; i < m; ++i) mu += s[i] * lambda[i];
            mu /= static_cast<T>(m);

            // convergence check
            T normR = T{0};
            for (auto r : rDual) normR += r * r;
            for (auto r : rPrimal) normR += r * r;
            if (std::sqrt(normR) < tol && mu < tol) {
                T obj = T{0};
                for (size_t j = 0; j < n; ++j) {
                    obj += c[j] * x[j];
                    for (size_t k = 0; k < n; ++k)
                        obj += T{0.5} * Q(j, k) * x[j] * x[k];
                }
                return { std::move(x), obj, iter + 1 };
            }

            // centering parameter σ
            T sigma = T{0.3};

            // Newton step (simplified: diagonal approximation)
            for (size_t j = 0; j < n; ++j) {
                T hess = Q(j, j);
                for (size_t i = 0; i < m; ++i)
                    hess += A(i, j) * A(i, j) * lambda[i] / s[i];
                if (hess > T{1e-15})
                    x[j] -= rDual[j] / hess * T{0.9};
            }

            // update s, λ
            for (size_t i = 0; i < m; ++i) {
                T Ax_i = T{0};
                for (size_t j = 0; j < n; ++j) Ax_i += A(i, j) * x[j];
                s[i] = std::max(b[i] - Ax_i, T{1e-10});
                lambda[i] = sigma * mu / s[i];
            }
        }

        T obj = T{0};
        for (size_t j = 0; j < n; ++j) {
            obj += c[j] * x[j];
            for (size_t k = 0; k < n; ++k)
                obj += T{0.5} * Q(j, k) * x[j] * x[k];
        }
        return { std::move(x), obj, maxIter };
    }

} // namespace sangi

#endif // SANGI_MATH_OPTIMIZATION_CONVEX_IMPL_HPP