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

// KMeans.hpp — k-means クラスタリング
//
// 移植元: lib++20 機械学習/クラスタリング/k-平均法.h
// 改良点:
//   - calx Matrix<T>/Vector<T> ベース
//   - concepts による型制約
//   - k-means++ 初期化
//   - テンプレート化 (float/double)

#ifndef CALX_ML_CLUSTERING_KMEANS_HPP
#define CALX_ML_CLUSTERING_KMEANS_HPP

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

namespace calx::ml {

    /// k-means クラスタリング結果
    template<typename T>
    struct KMeansResult {
        Matrix<T> centroids;            // (k × dim) 重心行列
        std::vector<int> labels;        // 各サンプルの所属クラスタ (0-based)
        int iterations;                 // 収束に要した反復回数
    };

    /// サンプル間のユークリッド距離の二乗
    template<typename T>
    static T squaredDistance(const Matrix<T>& data, std::size_t i,
                             const Matrix<T>& centroids, std::size_t c) {
        T sum = T{0};
        std::size_t dim = data.cols();
        for (std::size_t d = 0; d < dim; ++d) {
            T diff = data(i, d) - centroids(c, d);
            sum += diff * diff;
        }
        return sum;
    }

    /// k-means++ 初期化
    template<typename T>
    static Matrix<T> kmeansppInit(const Matrix<T>& data, std::size_t k,
                                   std::mt19937& rng) {
        std::size_t n = data.rows();
        std::size_t dim = data.cols();
        Matrix<T> centroids(k, dim);

        // 最初の重心をランダムに選択
        std::uniform_int_distribution<std::size_t> dist(0, n - 1);
        std::size_t first = dist(rng);
        for (std::size_t d = 0; d < dim; ++d)
            centroids(0, d) = data(first, d);

        // 残りの重心を距離に比例した確率で選択
        std::vector<T> minDist(n, std::numeric_limits<T>::max());

        for (std::size_t c = 1; c < k; ++c) {
            // 各点から最近の重心までの距離を更新
            for (std::size_t i = 0; i < n; ++i) {
                T d2 = squaredDistance(data, i, centroids, c - 1);
                if (d2 < minDist[i]) minDist[i] = d2;
            }

            // 距離の二乗に比例した確率で次の重心を選択
            std::discrete_distribution<std::size_t> wdist(minDist.begin(), minDist.end());
            std::size_t next = wdist(rng);
            for (std::size_t d = 0; d < dim; ++d)
                centroids(c, d) = data(next, d);
        }

        return centroids;
    }

    /// 各サンプルを最近の重心に割り当て
    template<typename T>
    static int assignCluster(const Matrix<T>& data, std::size_t i,
                              const Matrix<T>& centroids, std::size_t k) {
        T best = std::numeric_limits<T>::max();
        int label = 0;
        for (std::size_t c = 0; c < k; ++c) {
            T d2 = squaredDistance(data, i, centroids, c);
            if (d2 < best) {
                best = d2;
                label = static_cast<int>(c);
            }
        }
        return label;
    }

    /// k-means クラスタリング
    /// @param data     (n × dim) サンプル行列 (各行がひとつのサンプル)
    /// @param k        クラスタ数
    /// @param maxIter  最大反復回数
    /// @param seed     乱数シード
    template<typename T>
    KMeansResult<T> kmeans(const Matrix<T>& data, std::size_t k,
                           int maxIter = 100, unsigned seed = 42) {
        std::size_t n = data.rows();
        std::size_t dim = data.cols();
        assert(k >= 2 && k <= n);

        std::mt19937 rng(seed);

        // k-means++ 初期化
        Matrix<T> centroids = kmeansppInit(data, k, rng);

        std::vector<int> labels(n, -1);
        std::vector<int> prevLabels(n, -2);
        int iter = 0;

        for (; iter < maxIter; ++iter) {
            // 割り当てステップ
            for (std::size_t i = 0; i < n; ++i) {
                labels[i] = assignCluster(data, i, centroids, k);
            }

            // 収束チェック
            if (labels == prevLabels) break;
            prevLabels = labels;

            // 重心更新ステップ
            Matrix<T> newCentroids(k, dim, T{0});
            std::vector<int> counts(k, 0);

            for (std::size_t i = 0; i < n; ++i) {
                int c = labels[i];
                counts[c]++;
                for (std::size_t d = 0; d < dim; ++d) {
                    newCentroids(c, d) += data(i, d);
                }
            }

            for (std::size_t c = 0; c < k; ++c) {
                if (counts[c] > 0) {
                    for (std::size_t d = 0; d < dim; ++d) {
                        newCentroids(c, d) /= counts[c];
                    }
                }
            }

            centroids = std::move(newCentroids);
        }

        return { std::move(centroids), std::move(labels), iter };
    }

    /// 学習済み重心を使って新しいサンプルを分類
    template<typename T>
    int kmeansPredict(const Vector<T>& sample, const Matrix<T>& centroids) {
        std::size_t k = centroids.rows();
        std::size_t dim = centroids.cols();
        T best = std::numeric_limits<T>::max();
        int label = 0;
        for (std::size_t c = 0; c < k; ++c) {
            T d2 = T{0};
            for (std::size_t d = 0; d < dim; ++d) {
                T diff = sample[d] - centroids(c, d);
                d2 += diff * diff;
            }
            if (d2 < best) {
                best = d2;
                label = static_cast<int>(c);
            }
        }
        return label;
    }

    /// Mini-batch k-means
    /// @param data       (n × dim) サンプル行列
    /// @param k          クラスタ数
    /// @param batchSize  ミニバッチサイズ
    /// @param maxIter    最大反復回数
    /// @param seed       乱数シード
    template<typename T>
    KMeansResult<T> miniBatchKmeans(const Matrix<T>& data, std::size_t k,
                                     std::size_t batchSize = 100,
                                     int maxIter = 100, unsigned seed = 42)
    {
        std::size_t n = data.rows();
        std::size_t dim = data.cols();
        assert(k >= 2 && k <= n);
        if (batchSize > n) batchSize = n;

        std::mt19937 rng(seed);
        Matrix<T> centroids = kmeansppInit(data, k, rng);
        std::vector<int> clusterCounts(k, 0);
        std::vector<std::size_t> indices(n);
        std::iota(indices.begin(), indices.end(), 0);

        for (int iter = 0; iter < maxIter; ++iter) {
            std::shuffle(indices.begin(), indices.end(), rng);

            for (std::size_t b = 0; b < batchSize; ++b) {
                std::size_t idx = indices[b];
                int nearest = assignCluster(data, idx, centroids, k);
                clusterCounts[nearest]++;

                // オンライン重心更新: η = 1 / count
                T eta = T{1} / static_cast<T>(clusterCounts[nearest]);
                for (std::size_t d = 0; d < dim; ++d)
                    centroids(nearest, d) += eta * (data(idx, d) - centroids(nearest, d));
            }
        }

        // 最終ラベル割り当て
        std::vector<int> labels(n);
        for (std::size_t i = 0; i < n; ++i)
            labels[i] = assignCluster(data, i, centroids, k);

        return { std::move(centroids), std::move(labels), maxIter };
    }

} // namespace calx::ml

#endif // CALX_ML_CLUSTERING_KMEANS_HPP
