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

// background_subtraction.hpp — 背景差分 (MOG2)
//
// 新規実装:
//   - 混合ガウスモデル (Mixture of Gaussians)
//   - 前景/背景分離
//   - 適応的背景更新

#ifndef CALX_IMAGE_CV_BACKGROUND_SUBTRACTION_HPP
#define CALX_IMAGE_CV_BACKGROUND_SUBTRACTION_HPP

#include <image/Image.hpp>
#include <algorithm>
#include <cmath>
#include <vector>

namespace calx::cv {

    /// MOG2 背景差分
    class BackgroundSubtractorMOG2 {
    public:
        /// @param numGaussians  混合数 (通常 3-5)
        /// @param threshold     前景判定の閾値 (マハラノビス距離の二乗)
        /// @param learningRate  学習率 (0-1)
        BackgroundSubtractorMOG2(int numGaussians = 3, double threshold = 16.0,
                                  double learningRate = 0.01)
            : K_(numGaussians), threshold_(threshold), alpha_(learningRate) {}

        /// フレームを処理して前景マスクを返す
        /// @param frame   入力画像 (1ch)
        /// @return        前景マスク (1ch, 1.0=前景, 0.0=背景)
        template<typename T>
        image::Image<T> apply(const image::Image<T>& frame) {
            int H = static_cast<int>(frame.height());
            int W = static_cast<int>(frame.width());

            // 初期化
            if (models_.empty()) {
                models_.resize(H * W);
                for (auto& m : models_) {
                    m.gaussians.resize(K_);
                    m.gaussians[0] = {0.5, 128.0, 400.0, 1.0};  // 初期ガウス分布
                    for (int k = 1; k < K_; ++k) {
                        m.gaussians[k] = {0.0, 0.0, 400.0, 0.5 / (K_ - 1)};
                    }
                }
            }

            image::Image<T> mask(H, W);

            for (int y = 0; y < H; ++y) {
                for (int x = 0; x < W; ++x) {
                    double pixel = static_cast<double>(frame(y, x));
                    if constexpr (std::is_floating_point_v<T>) {
                        pixel *= 255.0;  // [0,1] → [0,255]
                    }

                    auto& model = models_[y * W + x];
                    bool matched = false;
                    bool isForeground = true;

                    // 重みの降順でマッチング
                    double weightSum = 0;
                    for (int k = 0; k < K_; ++k) {
                        auto& g = model.gaussians[k];
                        double diff = pixel - g.mean;
                        double mahal = diff * diff / g.variance;

                        if (mahal < threshold_) {
                            // マッチ: パラメータ更新
                            matched = true;
                            g.weight = (1 - alpha_) * g.weight + alpha_;
                            double rho = alpha_ / g.weight;
                            g.mean = (1 - rho) * g.mean + rho * pixel;
                            g.variance = (1 - rho) * g.variance + rho * diff * diff;
                            g.variance = std::max(g.variance, 4.0);

                            // 背景モデルの上位 K-1 個に含まれるか
                            weightSum += g.weight;
                            if (weightSum > 0.7) isForeground = false;
                        } else {
                            g.weight = (1 - alpha_) * g.weight;
                        }
                        weightSum += g.weight;
                    }

                    // マッチしなかった場合: 最小重みのガウスを置き換え
                    if (!matched) {
                        int minK = 0;
                        for (int k = 1; k < K_; ++k) {
                            if (model.gaussians[k].weight < model.gaussians[minK].weight)
                                minK = k;
                        }
                        model.gaussians[minK] = {alpha_, pixel, 400.0, alpha_};
                        isForeground = true;
                    }

                    // 重み正規化
                    double totalWeight = 0;
                    for (int k = 0; k < K_; ++k) totalWeight += model.gaussians[k].weight;
                    if (totalWeight > 0) {
                        for (int k = 0; k < K_; ++k) model.gaussians[k].weight /= totalWeight;
                    }

                    mask(y, x) = isForeground ? T{1} : T{0};
                    if constexpr (std::is_same_v<T, uint8_t>) {
                        mask(y, x) = isForeground ? T{255} : T{0};
                    }
                }
            }

            return mask;
        }

    private:
        struct Gaussian {
            double weight = 0;
            double mean = 0;
            double variance = 400.0;
            double sortKey = 0;  // weight/sqrt(variance)
        };

        struct PixelModel {
            std::vector<Gaussian> gaussians;
        };

        int K_;
        double threshold_;
        double alpha_;
        std::vector<PixelModel> models_;
    };

} // namespace calx::cv

#endif // CALX_IMAGE_CV_BACKGROUND_SUBTRACTION_HPP
