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

// UMAP.hpp — Uniform Manifold Approximation and Projection
//
// 近傍グラフ構築 + ファジー集合の交差 + SGD 最適化。
// 参考: McInnes, Healy & Melville (2018)

#ifndef CALX_ML_MANIFOLD_UMAP_HPP
#define CALX_ML_MANIFOLD_UMAP_HPP

#include <math/core/matrix.hpp>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <limits>
#include <numeric>
#include <random>
#include <vector>

namespace calx::ml {

    /// UMAP 結果
    template<typename T>
    struct UmapResult {
        Matrix<T> embedding;  // (n × outDim)
    };

    namespace detail {

        /// k 近傍をブルートフォースで取得 (距離², インデックス)
        template<typename T>
        std::vector<std::vector<std::pair<T, size_t>>>
        computeKNN(const Matrix<T>& X, size_t k)
        {
            size_t n = X.rows();
            size_t d = X.cols();
            std::vector<std::vector<std::pair<T, size_t>>> knn(n);

            for (size_t i = 0; i < n; ++i) {
                // 全点との距離を計算
                std::vector<std::pair<T, size_t>> dists(n);
                for (size_t j = 0; j < n; ++j) {
                    T dist = T(0);
                    for (size_t f = 0; f < d; ++f) {
                        T diff = X(i, f) - X(j, f);
                        dist += diff * diff;
                    }
                    dists[j] = { dist, j };
                }
                // 自分を除外してソート
                std::sort(dists.begin(), dists.end());
                knn[i].reserve(k);
                for (auto& [dist, idx] : dists) {
                    if (idx == i) continue;
                    knn[i].push_back({ std::sqrt(dist), idx });  // 距離 (√)
                    if (knn[i].size() == k) break;
                }
            }
            return knn;
        }

        /// sigma の探索 (二分探索): sum_j exp(-(d(i,j) - rho_i) / sigma_i) = log2(k)
        template<typename T>
        T findSigma(const std::vector<std::pair<T, size_t>>& neighbors,
                     T rho, T targetLog2k)
        {
            T lo = T(1e-10), hi = T(1000);
            T sigma = T(1);

            for (int iter = 0; iter < 64; ++iter) {
                sigma = (lo + hi) / T(2);
                T sum = T(0);
                for (auto& [dist, idx] : neighbors) {
                    T d = dist - rho;
                    if (d < T(0)) d = T(0);
                    sum += std::exp(-d / sigma);
                }
                if (std::abs(sum - targetLog2k) < T(1e-5)) break;
                if (sum > targetLog2k) lo = sigma;
                else hi = sigma;
            }
            return sigma;
        }

        /// ファジーグラフ (エッジリスト) の構築
        template<typename T>
        struct Edge {
            size_t i, j;
            T weight;
        };

        template<typename T>
        std::vector<Edge<T>> buildFuzzyGraph(
            const std::vector<std::vector<std::pair<T, size_t>>>& knn, size_t n)
        {
            size_t k = knn[0].size();
            T targetLog2k = std::log2(static_cast<T>(k));

            // rho_i = 最近傍までの距離
            std::vector<T> rho(n);
            std::vector<T> sigma(n);
            for (size_t i = 0; i < n; ++i) {
                rho[i] = knn[i].empty() ? T(0) : knn[i][0].first;
                sigma[i] = findSigma(knn[i], rho[i], targetLog2k);
            }

            // 条件付き確率
            // p_j|i = exp(-(d(i,j) - rho_i) / sigma_i)
            // 対称化: p_ij = p_j|i + p_i|j - p_j|i * p_i|j
            // エッジリストを対称行列として構築
            std::vector<std::vector<T>> W(n, std::vector<T>(n, T(0)));
            for (size_t i = 0; i < n; ++i) {
                for (auto& [dist, j] : knn[i]) {
                    T d = dist - rho[i];
                    if (d < T(0)) d = T(0);
                    T p = std::exp(-d / sigma[i]);
                    W[i][j] = std::max(W[i][j], p);
                }
            }

            // 対称化
            std::vector<Edge<T>> edges;
            for (size_t i = 0; i < n; ++i) {
                for (size_t j = i + 1; j < n; ++j) {
                    T a = W[i][j], b = W[j][i];
                    T sym = a + b - a * b;
                    if (sym > T(1e-10))
                        edges.push_back({ i, j, sym });
                }
            }
            return edges;
        }

    } // namespace detail

    /// UMAP を実行
    /// @param X           入力 (n × d)
    /// @param outDim      出力次元 (通常 2)
    /// @param nNeighbors  近傍数 (通常 15)
    /// @param minDist     最小距離 (埋め込み空間でのクラスタ密度制御)
    /// @param nEpochs     SGD エポック数
    /// @param learningRate 学習率
    /// @param seed        乱数シード
    template<typename T>
    UmapResult<T> umap(const Matrix<T>& X,
                        size_t outDim = 2,
                        size_t nNeighbors = 15,
                        T minDist = T(0.1),
                        size_t nEpochs = 500,
                        T learningRate = T(1),
                        unsigned seed = 42)
    {
        size_t n = X.rows();
        assert(n > 1);
        nNeighbors = std::min(nNeighbors, n - 1);

        // --- 近傍グラフ ---
        auto knn = detail::computeKNN(X, nNeighbors);
        auto edges = detail::buildFuzzyGraph(knn, n);

        // --- a, b パラメータ (minDist → 曲線フィット近似) ---
        // Phi(d) = 1 / (1 + a * d^(2b)) ≈ 1 if d < minDist, 快速減衰 otherwise
        // 簡易近似: a ≈ 1.93 * minDist^(-0.34), b ≈ 0.79 (参考値)
        T a, b;
        if (minDist < T(0.001)) {
            a = T(1.93); b = T(0.79);
        } else {
            b = T(0.79);
            a = std::pow(minDist, T(-1.93 * 0.79)) ;
            // より正確な近似
            a = T(1) / std::pow(minDist, T(2) * b);
        }

        // --- 初期化 (スペクトル初期化の簡易版: ランダム) ---
        std::mt19937 rng(seed);
        std::uniform_real_distribution<T> uniform(T(-10), T(10));
        Matrix<T> Y(n, outDim);
        for (size_t i = 0; i < n; ++i)
            for (size_t o = 0; o < outDim; ++o)
                Y(i, o) = uniform(rng);

        // --- エッジごとのエポック割り当て ---
        // 重みが大きいエッジほど頻繁にサンプリング
        T maxWeight = T(0);
        for (auto& e : edges)
            if (e.weight > maxWeight) maxWeight = e.weight;

        std::vector<T> epochsPerSample(edges.size());
        std::vector<T> epochNext(edges.size());
        for (size_t e = 0; e < edges.size(); ++e) {
            // 重みが大きいエッジほど epochsPerSample が小さい (頻繁にサンプリング)
            epochsPerSample[e] = static_cast<T>(nEpochs) /
                                  (edges[e].weight / maxWeight * static_cast<T>(nEpochs));
            if (epochsPerSample[e] < T(1)) epochsPerSample[e] = T(1);
            epochNext[e] = epochsPerSample[e];
        }

        // --- SGD 最適化 ---
        constexpr size_t negSamples = 5;
        std::uniform_int_distribution<size_t> randNode(0, n - 1);

        for (size_t epoch = 0; epoch < nEpochs; ++epoch) {
            T alpha = learningRate * (T(1) - static_cast<T>(epoch) / static_cast<T>(nEpochs));
            if (alpha < T(0)) alpha = T(0);

            for (size_t e = 0; e < edges.size(); ++e) {
                if (epochNext[e] > static_cast<T>(epoch)) continue;

                size_t i = edges[e].i;
                size_t j = edges[e].j;

                // 引力 (近いエッジを引き寄せる)
                T dist2 = T(0);
                for (size_t o = 0; o < outDim; ++o) {
                    T diff = Y(i, o) - Y(j, o);
                    dist2 += diff * diff;
                }
                if (dist2 < T(1e-10)) dist2 = T(1e-10);
                T grad_coeff = T(-2) * a * b * std::pow(dist2, b - T(1))
                               / (T(1) + a * std::pow(dist2, b));

                // クリッピング
                constexpr T clipVal = T(4);
                for (size_t o = 0; o < outDim; ++o) {
                    T diff = Y(i, o) - Y(j, o);
                    T grad = grad_coeff * diff;
                    if (grad > clipVal) grad = clipVal;
                    if (grad < -clipVal) grad = -clipVal;
                    Y(i, o) += alpha * grad;
                    Y(j, o) -= alpha * grad;
                }

                // 反発力 (ランダム負例)
                for (size_t neg = 0; neg < negSamples; ++neg) {
                    size_t k = randNode(rng);
                    if (k == i) continue;

                    T ndist2 = T(0);
                    for (size_t o = 0; o < outDim; ++o) {
                        T diff = Y(i, o) - Y(k, o);
                        ndist2 += diff * diff;
                    }
                    if (ndist2 < T(1e-10)) ndist2 = T(1e-10);

                    T rep_coeff = T(2) * b
                                  / ((T(0.001) + ndist2) * (T(1) + a * std::pow(ndist2, b)));

                    for (size_t o = 0; o < outDim; ++o) {
                        T diff = Y(i, o) - Y(k, o);
                        T grad = rep_coeff * diff;
                        if (grad > clipVal) grad = clipVal;
                        if (grad < -clipVal) grad = -clipVal;
                        Y(i, o) += alpha * grad;
                    }
                }

                epochNext[e] += epochsPerSample[e];
            }
        }

        return { std::move(Y) };
    }

} // namespace calx::ml

#endif // CALX_ML_MANIFOLD_UMAP_HPP
