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

// IsolationForest.hpp — Isolation Forest (異常検知)
//
// ランダム分割木のアンサンブルで異常スコアを計算。
// 異常な点ほど木の浅い位置で孤立する性質を利用。
// 参考: Liu, Ting & Zhou (2008)

#ifndef CALX_ML_ANOMALY_ISOLATIONFOREST_HPP
#define CALX_ML_ANOMALY_ISOLATIONFOREST_HPP

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

namespace calx::ml {

    namespace detail {

        // =================================================================
        // Isolation Tree ノード
        // =================================================================
        template<typename T>
        struct ITreeNode {
            bool isLeaf = false;
            size_t size = 0;         // 葉のサンプル数
            size_t featureIndex = 0; // 分割特徴量
            T splitValue = T(0);     // 分割値
            std::unique_ptr<ITreeNode> left;
            std::unique_ptr<ITreeNode> right;
        };

        /// BST の平均パス長 c(n) の計算
        /// c(n) = 2*H(n-1) - 2*(n-1)/n, H(i) = ln(i) + 0.5772156649 (Euler-Mascheroni)
        template<typename T>
        T averagePathLength(size_t n) {
            if (n <= 1) return T(0);
            if (n == 2) return T(1);
            T hn = std::log(static_cast<T>(n - 1)) + T(0.5772156649);
            return T(2) * hn - T(2) * static_cast<T>(n - 1) / static_cast<T>(n);
        }

        /// Isolation Tree を構築
        template<typename T>
        std::unique_ptr<ITreeNode<T>> buildITree(
            const Matrix<T>& X, std::vector<size_t>& indices,
            size_t currentDepth, size_t maxDepth, std::mt19937& rng)
        {
            auto node = std::make_unique<ITreeNode<T>>();
            size_t n = indices.size();

            // 葉ノード条件
            if (currentDepth >= maxDepth || n <= 1) {
                node->isLeaf = true;
                node->size = n;
                return node;
            }

            size_t d = X.cols();
            std::uniform_int_distribution<size_t> featDist(0, d - 1);

            // ランダム特徴量を選択し、その範囲内でランダム分割
            // 有効な分割を見つけるまで試行
            constexpr int maxTries = 10;
            for (int trial = 0; trial < maxTries; ++trial) {
                size_t f = featDist(rng);

                T minVal = X(indices[0], f);
                T maxVal = minVal;
                for (size_t i = 1; i < n; ++i) {
                    T v = X(indices[i], f);
                    if (v < minVal) minVal = v;
                    if (v > maxVal) maxVal = v;
                }

                if (minVal == maxVal) continue;  // この特徴量では分割不能

                std::uniform_real_distribution<T> splitDist(minVal, maxVal);
                T splitVal = splitDist(rng);

                std::vector<size_t> leftIdx, rightIdx;
                leftIdx.reserve(n);
                rightIdx.reserve(n);
                for (auto idx : indices) {
                    if (X(idx, f) < splitVal)
                        leftIdx.push_back(idx);
                    else
                        rightIdx.push_back(idx);
                }

                // どちらかが空ならリトライ
                if (leftIdx.empty() || rightIdx.empty()) continue;

                node->featureIndex = f;
                node->splitValue = splitVal;
                node->left = buildITree(X, leftIdx, currentDepth + 1, maxDepth, rng);
                node->right = buildITree(X, rightIdx, currentDepth + 1, maxDepth, rng);
                return node;
            }

            // 分割不能 → 葉
            node->isLeaf = true;
            node->size = n;
            return node;
        }

        /// パス長を計算
        template<typename T>
        T pathLength(const ITreeNode<T>* node, const Matrix<T>& X,
                      size_t row, size_t currentDepth)
        {
            if (node->isLeaf) {
                return static_cast<T>(currentDepth) + averagePathLength<T>(node->size);
            }
            if (X(row, node->featureIndex) < node->splitValue)
                return pathLength(node->left.get(), X, row, currentDepth + 1);
            else
                return pathLength(node->right.get(), X, row, currentDepth + 1);
        }

    } // namespace detail

    /// Isolation Forest 異常検知器
    template<typename T>
    class IsolationForest {
    public:
        /// @param numTrees       木の数
        /// @param sampleSize     各木のサブサンプルサイズ (0 = min(256, n))
        /// @param seed           乱数シード
        explicit IsolationForest(size_t numTrees = 100,
                                  size_t sampleSize = 0,
                                  unsigned seed = 42)
            : numTrees_(numTrees), sampleSize_(sampleSize), seed_(seed) {}

        /// 学習
        void fit(const Matrix<T>& X) {
            size_t n = X.rows();
            sampleSize_ = (sampleSize_ == 0) ? std::min(size_t(256), n) : sampleSize_;
            maxDepth_ = static_cast<size_t>(std::ceil(std::log2(static_cast<double>(sampleSize_))));
            avgPathN_ = detail::averagePathLength<T>(sampleSize_);

            std::mt19937 rng(seed_);
            trees_.clear();
            trees_.reserve(numTrees_);

            std::vector<size_t> allIndices(n);
            std::iota(allIndices.begin(), allIndices.end(), 0);

            for (size_t t = 0; t < numTrees_; ++t) {
                // サブサンプリング (非復元)
                std::vector<size_t> perm = allIndices;
                std::shuffle(perm.begin(), perm.end(), rng);
                std::vector<size_t> sample(perm.begin(),
                                            perm.begin() + std::min(sampleSize_, n));

                trees_.push_back(detail::buildITree(X, sample, 0, maxDepth_, rng));
            }

            data_ = &X;
        }

        /// 異常スコア (0〜1, 1 に近いほど異常)
        T score(const Matrix<T>& X, size_t row) const {
            T meanPath = T(0);
            for (auto& tree : trees_)
                meanPath += detail::pathLength(tree.get(), X, row, 0);
            meanPath /= static_cast<T>(trees_.size());

            // s(x, n) = 2^(-E[h(x)] / c(n))
            return std::pow(T(2), -meanPath / avgPathN_);
        }

        /// 全サンプルの異常スコア
        std::vector<T> score(const Matrix<T>& X) const {
            std::vector<T> result(X.rows());
            for (size_t i = 0; i < X.rows(); ++i)
                result[i] = score(X, i);
            return result;
        }

        /// 異常判定 (スコア > threshold)
        std::vector<bool> predict(const Matrix<T>& X, T threshold = T(0.5)) const {
            auto scores = score(X);
            std::vector<bool> result(scores.size());
            for (size_t i = 0; i < scores.size(); ++i)
                result[i] = scores[i] > threshold;
            return result;
        }

        size_t numTrees() const { return trees_.size(); }

    private:
        size_t numTrees_;
        size_t sampleSize_;
        unsigned seed_;
        size_t maxDepth_ = 0;
        T avgPathN_ = T(0);
        std::vector<std::unique_ptr<detail::ITreeNode<T>>> trees_;
        const Matrix<T>* data_ = nullptr;
    };

} // namespace calx::ml

#endif // CALX_ML_ANOMALY_ISOLATIONFOREST_HPP
