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

// HMM.hpp — 隠れマルコフモデル + Viterbi アルゴリズム
//
// 移植元: lib++20 数学/最適化/多変数/Viterbi.h
// 改良点:
//   - calx Matrix<T>/Vector<T> ベース
//   - 対数確率演算 (アンダーフロー対策)
//   - Forward アルゴリズム追加

#ifndef CALX_ML_HMM_HPP
#define CALX_ML_HMM_HPP

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

namespace calx::ml {

    /// 隠れマルコフモデル (対数確率で保持)
    template<typename T>
    class HMM {
    public:
        int numStates;          // 状態数
        int numSymbols;         // 観測シンボル数
        Matrix<T> logTransition;   // 状態遷移確率 (対数)
        Matrix<T> logEmission;     // シンボル出力確率 (対数)
        Vector<T> logInitial;      // 初期状態確率 (対数)

        /// @param states      状態数
        /// @param symbols     観測シンボル数
        /// @param transition  状態遷移確率 (states × states, 生の確率)
        /// @param emission    シンボル出力確率 (states × symbols, 生の確率)
        /// @param initial     初期状態確率 (states, 生の確率)
        HMM(int states, int symbols,
            const Matrix<T>& transition,
            const Matrix<T>& emission,
            const Vector<T>& initial)
            : numStates(states)
            , numSymbols(symbols)
            , logTransition(states, states)
            , logEmission(states, symbols)
            , logInitial(states)
        {
            assert(states > 0 && symbols > 0);
            for (int i = 0; i < states; ++i) {
                logInitial[i] = std::log(initial[i]);
                for (int j = 0; j < states; ++j)
                    logTransition(i, j) = std::log(transition(i, j));
                for (int j = 0; j < symbols; ++j)
                    logEmission(i, j) = std::log(emission(i, j));
            }
        }
    };

    /// Viterbi アルゴリズム: 最尤状態系列を求める
    /// @param hmm          隠れマルコフモデル
    /// @param observations 観測シンボル列 (0-based インデックス)
    /// @return             最尤状態系列
    template<typename T>
    std::vector<int> viterbi(const HMM<T>& hmm, const std::vector<int>& observations) {
        int N = static_cast<int>(observations.size());
        assert(N > 0);

        int S = hmm.numStates;

        // prob(t, s) = 時刻 t で状態 s にいる最大対数尤度
        Matrix<T> prob(N, S, std::numeric_limits<T>::lowest());
        Matrix<int> prev(N, S, 0);

        // 初期化
        for (int s = 0; s < S; ++s) {
            prob(0, s) = hmm.logInitial[s] + hmm.logEmission(s, observations[0]);
        }

        // 再帰
        for (int t = 1; t < N; ++t) {
            for (int s = 0; s < S; ++s) {
                T best = std::numeric_limits<T>::lowest();
                int bestPrev = 0;
                for (int k = 0; k < S; ++k) {
                    T p = prob(t - 1, k) + hmm.logTransition(k, s);
                    if (p > best) {
                        best = p;
                        bestPrev = k;
                    }
                }
                prob(t, s) = hmm.logEmission(s, observations[t]) + best;
                prev(t - 1, s) = bestPrev;
            }
        }

        // バックトレース
        std::vector<int> path(N);
        T best = std::numeric_limits<T>::lowest();
        for (int s = 0; s < S; ++s) {
            if (prob(N - 1, s) > best) {
                best = prob(N - 1, s);
                path[N - 1] = s;
            }
        }
        for (int t = N - 2; t >= 0; --t) {
            path[t] = prev(t, path[t + 1]);
        }

        return path;
    }

    /// Forward アルゴリズム: 観測系列の尤度 (対数) を求める
    template<typename T>
    T forward(const HMM<T>& hmm, const std::vector<int>& observations) {
        int N = static_cast<int>(observations.size());
        assert(N > 0);

        int S = hmm.numStates;
        Matrix<T> alpha(N, S);

        // 初期化
        for (int s = 0; s < S; ++s) {
            alpha(0, s) = hmm.logInitial[s] + hmm.logEmission(s, observations[0]);
        }

        // 再帰 (log-sum-exp で精度を保つ)
        for (int t = 1; t < N; ++t) {
            for (int s = 0; s < S; ++s) {
                T maxVal = std::numeric_limits<T>::lowest();
                for (int k = 0; k < S; ++k) {
                    T v = alpha(t - 1, k) + hmm.logTransition(k, s);
                    if (v > maxVal) maxVal = v;
                }
                T sum = T{0};
                for (int k = 0; k < S; ++k) {
                    sum += std::exp(alpha(t - 1, k) + hmm.logTransition(k, s) - maxVal);
                }
                alpha(t, s) = hmm.logEmission(s, observations[t]) + maxVal + std::log(sum);
            }
        }

        // 最終尤度 (log-sum-exp)
        T maxVal = std::numeric_limits<T>::lowest();
        for (int s = 0; s < S; ++s)
            if (alpha(N - 1, s) > maxVal) maxVal = alpha(N - 1, s);

        T sum = T{0};
        for (int s = 0; s < S; ++s)
            sum += std::exp(alpha(N - 1, s) - maxVal);

        return maxVal + std::log(sum);
    }

    /// Baum-Welch アルゴリズム (HMM パラメータ推定)
    /// @param observations  観測系列
    /// @param numStates     状態数
    /// @param numSymbols    観測シンボル数
    /// @param maxIter       最大反復回数
    /// @param tol           対数尤度の収束閾値
    /// @param seed          乱数シード
    /// @return 学習済み HMM
    template<typename T>
    HMM<T> baumWelch(const std::vector<int>& observations,
                      int numStates, int numSymbols,
                      int maxIter = 100, T tol = T{1e-6}, unsigned seed = 42)
    {
        int S = numStates, M = numSymbols, N = static_cast<int>(observations.size());
        std::mt19937 rng(seed);

        // ランダム初期化 (行和=1 に正規化)
        Matrix<T> trans(S, S);
        Matrix<T> emit(S, M);
        Vector<T> init(S);
        {
            std::uniform_real_distribution<T> dist(T{0.1}, T{1});
            for (int i = 0; i < S; ++i) {
                T sum = T{0};
                for (int j = 0; j < S; ++j) { trans(i, j) = dist(rng); sum += trans(i, j); }
                for (int j = 0; j < S; ++j) trans(i, j) /= sum;
                sum = T{0};
                for (int j = 0; j < M; ++j) { emit(i, j) = dist(rng); sum += emit(i, j); }
                for (int j = 0; j < M; ++j) emit(i, j) /= sum;
            }
            T sum = T{0};
            for (int i = 0; i < S; ++i) { init[i] = dist(rng); sum += init[i]; }
            for (int i = 0; i < S; ++i) init[i] /= sum;
        }

        T prevLL = std::numeric_limits<T>::lowest();

        for (int iter = 0; iter < maxIter; ++iter) {
            // Forward
            Matrix<T> alpha(N, S);
            for (int s = 0; s < S; ++s)
                alpha(0, s) = init[s] * emit(s, observations[0]);
            for (int t = 1; t < N; ++t)
                for (int s = 0; s < S; ++s) {
                    T sum = T{0};
                    for (int k = 0; k < S; ++k) sum += alpha(t - 1, k) * trans(k, s);
                    alpha(t, s) = sum * emit(s, observations[t]);
                }

            // Backward
            Matrix<T> beta(N, S);
            for (int s = 0; s < S; ++s) beta(N - 1, s) = T{1};
            for (int t = N - 2; t >= 0; --t)
                for (int s = 0; s < S; ++s) {
                    T sum = T{0};
                    for (int k = 0; k < S; ++k)
                        sum += trans(s, k) * emit(k, observations[t + 1]) * beta(t + 1, k);
                    beta(t, s) = sum;
                }

            // 対数尤度
            T ll = T{0};
            for (int s = 0; s < S; ++s) ll += alpha(N - 1, s);
            T logLL = std::log(std::max(ll, T{1e-300}));
            if (logLL - prevLL < tol && iter > 0) break;
            prevLL = logLL;

            // Gamma, Xi
            Matrix<T> gamma(N, S);
            for (int t = 0; t < N; ++t) {
                T sum = T{0};
                for (int s = 0; s < S; ++s) sum += alpha(t, s) * beta(t, s);
                if (sum < T{1e-300}) sum = T{1e-300};
                for (int s = 0; s < S; ++s) gamma(t, s) = alpha(t, s) * beta(t, s) / sum;
            }

            // パラメータ更新
            // 初期確率
            for (int s = 0; s < S; ++s) init[s] = gamma(0, s);

            // 遷移確率
            for (int i = 0; i < S; ++i) {
                T gammaSum = T{0};
                for (int t = 0; t < N - 1; ++t) gammaSum += gamma(t, i);
                if (gammaSum < T{1e-300}) gammaSum = T{1e-300};

                for (int j = 0; j < S; ++j) {
                    T xiSum = T{0};
                    for (int t = 0; t < N - 1; ++t) {
                        T denom = T{0};
                        for (int k = 0; k < S; ++k) denom += alpha(t, k) * beta(t, k);
                        if (denom < T{1e-300}) denom = T{1e-300};
                        xiSum += alpha(t, i) * trans(i, j) * emit(j, observations[t + 1])
                                 * beta(t + 1, j) / denom;
                    }
                    trans(i, j) = xiSum / gammaSum;
                }
            }

            // 放射確率
            for (int s = 0; s < S; ++s) {
                T gammaSum = T{0};
                for (int t = 0; t < N; ++t) gammaSum += gamma(t, s);
                if (gammaSum < T{1e-300}) gammaSum = T{1e-300};

                for (int m = 0; m < M; ++m) {
                    T sum = T{0};
                    for (int t = 0; t < N; ++t)
                        if (observations[t] == m) sum += gamma(t, s);
                    emit(s, m) = sum / gammaSum;
                }
            }
        }

        return HMM<T>(numStates, numSymbols, trans, emit, init);
    }

} // namespace calx::ml

#endif // CALX_ML_HMM_HPP
