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

// HierarchicalClustering.hpp — 階層的クラスタリング
//
// 新規実装 (scipy.cluster.hierarchy 対応)
// Ward 法, 単連結法, 完全連結法, 群平均法

#ifndef CALX_ML_CLUSTERING_HIERARCHICAL_HPP
#define CALX_ML_CLUSTERING_HIERARCHICAL_HPP

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

namespace calx::ml {

    /// 連結法
    enum class Linkage {
        Single,     // 単連結法 (最短距離)
        Complete,   // 完全連結法 (最長距離)
        Average,    // 群平均法
        Ward        // Ward 法 (分散最小化)
    };

    /// デンドログラムの 1 ステップ
    struct Merge {
        int cluster1;       // 統合されたクラスタ 1 (0-based)
        int cluster2;       // 統合されたクラスタ 2
        double distance;    // 統合時の距離
        int size;           // 統合後のクラスタサイズ
    };

    /// 距離行列を計算 (condensed form: n*(n-1)/2)
    template<typename T>
    static std::vector<T> computeDistanceMatrix(const Matrix<T>& data) {
        std::size_t n = data.rows();
        std::size_t dim = data.cols();
        std::vector<T> dist(n * (n - 1) / 2);

        std::size_t idx = 0;
        for (std::size_t i = 0; i < n; ++i) {
            for (std::size_t j = i + 1; j < n; ++j) {
                T sum = T{0};
                for (std::size_t d = 0; d < dim; ++d) {
                    T diff = data(i, d) - data(j, d);
                    sum += diff * diff;
                }
                dist[idx++] = std::sqrt(sum);
            }
        }
        return dist;
    }

    /// condensed index: (i, j) → idx (i < j)
    inline std::size_t condensedIndex(std::size_t n, std::size_t i, std::size_t j) {
        if (i > j) std::swap(i, j);
        return n * i - i * (i + 1) / 2 + j - i - 1;
    }

    /// 階層的クラスタリング
    /// @param data     (n × dim) サンプル行列
    /// @param method   連結法
    /// @return         デンドログラム (n-1 ステップ)
    template<typename T>
    std::vector<Merge> hierarchicalClustering(const Matrix<T>& data, Linkage method = Linkage::Ward) {
        std::size_t n = data.rows();
        assert(n >= 2);

        // クラスタ間距離行列 (full n×n, 上三角を使用)
        // Ward 法では重心を管理する必要がある
        std::vector<T> dist = computeDistanceMatrix(data);

        // 各クラスタのサイズ
        std::vector<int> sizes(n, 1);
        // アクティブなクラスタ
        std::vector<bool> active(n, true);
        // Ward 法用: 各クラスタの重心
        std::size_t dim = data.cols();
        Matrix<T> centroids(n, dim);
        for (std::size_t i = 0; i < n; ++i)
            for (std::size_t d = 0; d < dim; ++d)
                centroids(i, d) = data(i, d);

        std::vector<Merge> dendrogram;
        dendrogram.reserve(n - 1);

        // 距離参照 (condensed → i,j)
        auto getDist = [&](std::size_t i, std::size_t j) -> T {
            if (i == j) return T{0};
            return dist[condensedIndex(n, i, j)];
        };
        auto setDist = [&](std::size_t i, std::size_t j, T val) {
            if (i != j) dist[condensedIndex(n, i, j)] = val;
        };

        for (std::size_t step = 0; step < n - 1; ++step) {
            // 最近のペアを探す
            T minDist = std::numeric_limits<T>::max();
            std::size_t ci = 0, cj = 0;

            for (std::size_t i = 0; i < n; ++i) {
                if (!active[i]) continue;
                for (std::size_t j = i + 1; j < n; ++j) {
                    if (!active[j]) continue;
                    T d;
                    if (method == Linkage::Ward) {
                        // Ward 距離: Δ = (ni*nj)/(ni+nj) * ||ci - cj||^2
                        T sum = T{0};
                        for (std::size_t dd = 0; dd < dim; ++dd) {
                            T diff = centroids(i, dd) - centroids(j, dd);
                            sum += diff * diff;
                        }
                        d = static_cast<T>(sizes[i]) * static_cast<T>(sizes[j])
                            / (static_cast<T>(sizes[i]) + static_cast<T>(sizes[j])) * sum;
                    } else {
                        d = getDist(i, j);
                    }
                    if (d < minDist) {
                        minDist = d;
                        ci = i; cj = j;
                    }
                }
            }

            // ci に cj を統合
            dendrogram.push_back({
                static_cast<int>(ci), static_cast<int>(cj),
                static_cast<double>(std::sqrt(minDist)),
                sizes[ci] + sizes[cj]
            });

            // ci と他クラスタの距離を更新
            for (std::size_t k = 0; k < n; ++k) {
                if (!active[k] || k == ci || k == cj) continue;

                T dik = getDist(ci, k);
                T djk = getDist(cj, k);
                T newDist;

                switch (method) {
                    case Linkage::Single:
                        newDist = std::min(dik, djk);
                        break;
                    case Linkage::Complete:
                        newDist = std::max(dik, djk);
                        break;
                    case Linkage::Average:
                        newDist = (dik * sizes[ci] + djk * sizes[cj])
                                  / (sizes[ci] + sizes[cj]);
                        break;
                    case Linkage::Ward:
                        // Lance-Williams for Ward
                        newDist = dik; // 距離は次の反復で再計算
                        break;
                }
                setDist(ci, k, newDist);
            }

            // Ward 法: 重心を更新
            if (method == Linkage::Ward) {
                T totalSize = static_cast<T>(sizes[ci] + sizes[cj]);
                for (std::size_t d = 0; d < dim; ++d) {
                    centroids(ci, d) = (centroids(ci, d) * sizes[ci]
                                        + centroids(cj, d) * sizes[cj]) / totalSize;
                }
            }

            sizes[ci] += sizes[cj];
            active[cj] = false;
        }

        return dendrogram;
    }

    /// デンドログラムをカットして k クラスタに分割
    /// @param dendrogram   階層的クラスタリングの結果
    /// @param n            元のサンプル数
    /// @param k            分割するクラスタ数
    /// @return             各サンプルのクラスタ番号 (0-based)
    inline std::vector<int> cutDendrogram(const std::vector<Merge>& dendrogram,
                                           std::size_t n, int k) {
        assert(k >= 1 && static_cast<std::size_t>(k) <= n);

        // Union-Find
        std::vector<int> parent(n);
        std::iota(parent.begin(), parent.end(), 0);

        std::function<int(int)> find = [&](int x) -> int {
            return parent[x] == x ? x : (parent[x] = find(parent[x]));
        };

        // 最後の k-1 ステップを除いた統合を適用
        std::size_t steps = n - k;
        for (std::size_t i = 0; i < steps && i < dendrogram.size(); ++i) {
            int a = find(dendrogram[i].cluster1);
            int b = find(dendrogram[i].cluster2);
            parent[b] = a;
        }

        // ラベルを 0-based に再番号付け
        std::vector<int> labels(n);
        std::unordered_map<int, int> labelMap;
        int nextLabel = 0;
        for (std::size_t i = 0; i < n; ++i) {
            int root = find(static_cast<int>(i));
            if (labelMap.find(root) == labelMap.end()) {
                labelMap[root] = nextLabel++;
            }
            labels[i] = labelMap[root];
        }
        return labels;
    }

} // namespace calx::ml

#endif // CALX_ML_CLUSTERING_HIERARCHICAL_HPP
