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

// TSNE.hpp — t-SNE (t-distributed Stochastic Neighbor Embedding)
//
// 高次元データの 2D/3D 可視化用次元削減。
// Barnes-Hut 近似 (quad-tree) による O(N log N) 実装。
// 参考: van der Maaten & Hinton (2008), van der Maaten (2014)

#ifndef CALX_ML_MANIFOLD_TSNE_HPP
#define CALX_ML_MANIFOLD_TSNE_HPP

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

namespace calx::ml {

    namespace detail {

        // =================================================================
        // Quad-tree (Barnes-Hut 用)
        // =================================================================
        template<typename T>
        struct QuadTreeNode {
            // 境界ボックス
            T cx, cy, hw, hh;  // center x/y, half-width/height
            // 質量中心
            T comX = T(0), comY = T(0);
            size_t count = 0;
            // 子ノード (NW, NE, SW, SE)
            std::unique_ptr<QuadTreeNode> children[4];
            bool isLeaf = true;
            // 単一点 (リーフ)
            T px = T(0), py = T(0);

            QuadTreeNode(T cx_, T cy_, T hw_, T hh_)
                : cx(cx_), cy(cy_), hw(hw_), hh(hh_) {}

            int quadrant(T x, T y) const {
                bool west = x <= cx;
                bool north = y <= cy;
                if (north && west) return 0;   // NW
                if (north && !west) return 1;   // NE
                if (!north && west) return 2;   // SW
                return 3;                        // SE
            }

            void subdivide() {
                T qw = hw / T(2), qh = hh / T(2);
                children[0] = std::make_unique<QuadTreeNode>(cx - qw, cy - qh, qw, qh);
                children[1] = std::make_unique<QuadTreeNode>(cx + qw, cy - qh, qw, qh);
                children[2] = std::make_unique<QuadTreeNode>(cx - qw, cy + qh, qw, qh);
                children[3] = std::make_unique<QuadTreeNode>(cx + qw, cy + qh, qw, qh);
                isLeaf = false;
            }

            void insert(T x, T y) {
                if (count == 0) {
                    px = x; py = y;
                    comX = x; comY = y;
                    count = 1;
                    return;
                }

                // 質量中心の更新
                T newCount = static_cast<T>(count + 1);
                comX = (comX * static_cast<T>(count) + x) / newCount;
                comY = (comY * static_cast<T>(count) + y) / newCount;

                if (isLeaf) {
                    subdivide();
                    // 既存の点を子に移動
                    children[quadrant(px, py)]->insert(px, py);
                }

                children[quadrant(x, y)]->insert(x, y);
                count++;
            }

            /// Barnes-Hut 反発力の計算
            /// F = (1 / (1 + ||y_i - y_j||^2)) の方向成分を累積
            void computeRepulsion(T yi0, T yi1, T theta,
                                   T& fRep0, T& fRep1, T& zSum) const
            {
                if (count == 0) return;

                T dx = yi0 - comX;
                T dy = yi1 - comY;
                T dist2 = dx * dx + dy * dy;

                // 十分遠ければ質量中心で近似
                T cellSize = hw * T(2);  // セルの幅
                if (isLeaf || (cellSize * cellSize < theta * theta * dist2)) {
                    T qij = T(1) / (T(1) + dist2);
                    T mult = static_cast<T>(count) * qij * qij;
                    fRep0 += mult * dx;
                    fRep1 += mult * dy;
                    zSum += static_cast<T>(count) * qij;
                    return;
                }

                for (int c = 0; c < 4; ++c) {
                    if (children[c])
                        children[c]->computeRepulsion(yi0, yi1, theta,
                                                       fRep0, fRep1, zSum);
                }
            }
        };

        // =================================================================
        // perplexity → sigma 探索 (二分探索)
        // =================================================================
        template<typename T>
        T computePerplexityRow(const std::vector<T>& dist2, size_t i, T beta,
                                std::vector<T>& P_row)
        {
            size_t n = dist2.size();
            T sumP = T(0);
            for (size_t j = 0; j < n; ++j) {
                if (j == i) { P_row[j] = T(0); continue; }
                P_row[j] = std::exp(-dist2[j] * beta);
                sumP += P_row[j];
            }
            if (sumP < T(1e-100)) sumP = T(1e-100);

            T H = T(0);
            for (size_t j = 0; j < n; ++j) {
                P_row[j] /= sumP;
                if (P_row[j] > T(1e-100))
                    H -= P_row[j] * std::log2(P_row[j]);
            }
            return H;  // log2 perplexity
        }

        template<typename T>
        void findBeta(const std::vector<T>& dist2, size_t i, T targetH,
                       std::vector<T>& P_row)
        {
            T beta = T(1);
            T betaMin = T(0);
            T betaMax = T(1e10);
            constexpr int maxBinSearch = 50;

            for (int iter = 0; iter < maxBinSearch; ++iter) {
                T H = computePerplexityRow(dist2, i, beta, P_row);
                T diff = H - targetH;
                if (std::abs(diff) < T(1e-5)) break;

                if (diff > T(0)) {
                    betaMin = beta;
                    beta = (betaMax >= T(1e9)) ? beta * T(2) : (beta + betaMax) / T(2);
                } else {
                    betaMax = beta;
                    beta = (beta + betaMin) / T(2);
                }
            }
            computePerplexityRow(dist2, i, beta, P_row);
        }

    } // namespace detail

    /// t-SNE 結果
    template<typename T>
    struct TsneResult {
        Matrix<T> embedding;  // (n × outDim)
        T finalCost;          // 最終 KL ダイバージェンス
    };

    /// t-SNE を実行
    /// @param X           入力 (n × d)
    /// @param outDim      出力次元 (通常 2)
    /// @param perplexity  パープレキシティ (通常 5-50)
    /// @param maxIter     最大反復回数
    /// @param learningRate 学習率
    /// @param theta       Barnes-Hut 近似パラメータ (0 = 厳密, 0.5 が標準)
    /// @param seed        乱数シード
    template<typename T>
    TsneResult<T> tsne(const Matrix<T>& X,
                        size_t outDim = 2,
                        T perplexity = T(30),
                        size_t maxIter = 1000,
                        T learningRate = T(200),
                        T theta = T(0.5),
                        unsigned seed = 42)
    {
        size_t n = X.rows();
        size_t d = X.cols();
        assert(n > 1);

        // --- 距離行列 & 条件付き確率 P ---
        T targetH = std::log2(perplexity);
        std::vector<std::vector<T>> P(n, std::vector<T>(n, T(0)));

        // ペアワイズ距離
        std::vector<T> dist2Row(n);
        for (size_t i = 0; i < n; ++i) {
            for (size_t j = 0; j < n; ++j) {
                T s = T(0);
                for (size_t f = 0; f < d; ++f) {
                    T diff = X(i, f) - X(j, f);
                    s += diff * diff;
                }
                dist2Row[j] = s;
            }
            detail::findBeta(dist2Row, i, targetH, P[i]);
        }

        // 対称化: P_ij = (P_i|j + P_j|i) / (2n)
        for (size_t i = 0; i < n; ++i)
            for (size_t j = i + 1; j < n; ++j) {
                T sym = (P[i][j] + P[j][i]) / static_cast<T>(2 * n);
                P[i][j] = sym;
                P[j][i] = sym;
            }

        // Early exaggeration: P *= 4
        for (size_t i = 0; i < n; ++i)
            for (size_t j = 0; j < n; ++j)
                P[i][j] *= T(4);

        // --- 初期化 ---
        std::mt19937 rng(seed);
        std::normal_distribution<T> normal(T(0), T(1e-4));
        Matrix<T> Y(n, outDim);
        for (size_t i = 0; i < n; ++i)
            for (size_t o = 0; o < outDim; ++o)
                Y(i, o) = normal(rng);

        // モメンタム
        Matrix<T> gains(n, outDim, T(1));
        Matrix<T> dY(n, outDim, T(0));
        Matrix<T> uY(n, outDim, T(0));  // velocity

        constexpr T momentum_init = T(0.5);
        constexpr T momentum_final = T(0.8);
        constexpr size_t exaggeration_end = 250;

        T finalCost = T(0);

        for (size_t iter = 0; iter < maxIter; ++iter) {
            T mom = (iter < 250) ? momentum_init : momentum_final;

            if (iter == exaggeration_end) {
                // Early exaggeration 終了
                for (size_t i = 0; i < n; ++i)
                    for (size_t j = 0; j < n; ++j)
                        P[i][j] /= T(4);
            }

            // 勾配計算
            if (theta > T(0) && outDim == 2) {
                // --- Barnes-Hut ---
                // 境界ボックス
                T minX = Y(0, 0), maxX = Y(0, 0);
                T minY = Y(0, 1), maxY = Y(0, 1);
                for (size_t i = 1; i < n; ++i) {
                    if (Y(i, 0) < minX) minX = Y(i, 0);
                    if (Y(i, 0) > maxX) maxX = Y(i, 0);
                    if (Y(i, 1) < minY) minY = Y(i, 1);
                    if (Y(i, 1) > maxY) maxY = Y(i, 1);
                }
                T cx = (minX + maxX) / T(2), cy = (minY + maxY) / T(2);
                T hw = (maxX - minX) / T(2) + T(1e-5);
                T hh = (maxY - minY) / T(2) + T(1e-5);
                if (hw < hh) hw = hh; else hh = hw;

                detail::QuadTreeNode<T> root(cx, cy, hw, hh);
                for (size_t i = 0; i < n; ++i)
                    root.insert(Y(i, 0), Y(i, 1));

                // 引力 + 反発力
                for (size_t i = 0; i < n; ++i) {
                    // 引力 (P > 0 の近傍のみ)
                    T fAttr0 = T(0), fAttr1 = T(0);
                    for (size_t j = 0; j < n; ++j) {
                        if (i == j || P[i][j] < T(1e-100)) continue;
                        T dx = Y(i, 0) - Y(j, 0);
                        T dy = Y(i, 1) - Y(j, 1);
                        T qij = T(1) / (T(1) + dx * dx + dy * dy);
                        fAttr0 += P[i][j] * qij * dx;
                        fAttr1 += P[i][j] * qij * dy;
                    }

                    // 反発力 (Barnes-Hut)
                    T fRep0 = T(0), fRep1 = T(0), zSum = T(0);
                    root.computeRepulsion(Y(i, 0), Y(i, 1), theta,
                                           fRep0, fRep1, zSum);

                    if (zSum < T(1e-100)) zSum = T(1e-100);
                    dY(i, 0) = T(4) * (fAttr0 - fRep0 / zSum);
                    dY(i, 1) = T(4) * (fAttr1 - fRep1 / zSum);
                }
            } else {
                // --- 厳密計算 ---
                // Q 行列
                T Z = T(0);
                std::vector<std::vector<T>> Q(n, std::vector<T>(n, T(0)));
                for (size_t i = 0; i < n; ++i)
                    for (size_t j = i + 1; j < n; ++j) {
                        T dist = T(0);
                        for (size_t o = 0; o < outDim; ++o) {
                            T diff = Y(i, o) - Y(j, o);
                            dist += diff * diff;
                        }
                        T q = T(1) / (T(1) + dist);
                        Q[i][j] = q;
                        Q[j][i] = q;
                        Z += T(2) * q;
                    }
                if (Z < T(1e-100)) Z = T(1e-100);

                for (size_t i = 0; i < n; ++i) {
                    for (size_t o = 0; o < outDim; ++o) {
                        T grad = T(0);
                        for (size_t j = 0; j < n; ++j) {
                            if (i == j) continue;
                            T pq = P[i][j] - Q[i][j] / Z;
                            grad += pq * Q[i][j] * (Y(i, o) - Y(j, o));
                        }
                        dY(i, o) = T(4) * grad;
                    }
                }
            }

            // 更新 (adaptive gains + momentum)
            for (size_t i = 0; i < n; ++i)
                for (size_t o = 0; o < outDim; ++o) {
                    bool sameSign = (dY(i, o) > T(0)) == (uY(i, o) > T(0));
                    gains(i, o) = sameSign
                        ? std::max(gains(i, o) * T(0.8), T(0.01))
                        : gains(i, o) + T(0.2);

                    uY(i, o) = mom * uY(i, o) - learningRate * gains(i, o) * dY(i, o);
                    Y(i, o) += uY(i, o);
                }

            // 中心化
            for (size_t o = 0; o < outDim; ++o) {
                T mean = T(0);
                for (size_t i = 0; i < n; ++i) mean += Y(i, o);
                mean /= static_cast<T>(n);
                for (size_t i = 0; i < n; ++i) Y(i, o) -= mean;
            }
        }

        // KL ダイバージェンス
        finalCost = T(0);
        T Z = T(0);
        for (size_t i = 0; i < n; ++i)
            for (size_t j = i + 1; j < n; ++j) {
                T dist = T(0);
                for (size_t o = 0; o < outDim; ++o) {
                    T diff = Y(i, o) - Y(j, o);
                    dist += diff * diff;
                }
                Z += T(2) / (T(1) + dist);
            }
        if (Z < T(1e-100)) Z = T(1e-100);
        for (size_t i = 0; i < n; ++i)
            for (size_t j = 0; j < n; ++j) {
                if (i == j || P[i][j] < T(1e-100)) continue;
                T dist = T(0);
                for (size_t o = 0; o < outDim; ++o) {
                    T diff = Y(i, o) - Y(j, o);
                    dist += diff * diff;
                }
                T qij = (T(1) / (T(1) + dist)) / Z;
                if (qij < T(1e-100)) qij = T(1e-100);
                finalCost += P[i][j] * std::log(P[i][j] / qij);
            }

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

} // namespace calx::ml

#endif // CALX_ML_MANIFOLD_TSNE_HPP
