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

/**
 * @file denoise.hpp
 * @brief Non-Local Means デノイズ (OCV-24)
 */

#ifndef CALX_IMAGE_IMGPROC_DENOISE_HPP
#define CALX_IMAGE_IMGPROC_DENOISE_HPP

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

namespace calx::image {

/// Non-Local Means デノイズ
/// h: フィルタ強度, templateSize: パッチサイズ, searchSize: 探索窓サイズ
template<typename T>
Image<T> nlMeansFilter(const Image<T>& src, double h = 10.0,
                        int templateSize = 7, int searchSize = 21) {
    int H = static_cast<int>(src.height());
    int W = static_cast<int>(src.width());
    int ch = src.numChannels();
    int tHalf = templateSize / 2;
    int sHalf = searchSize / 2;
    double h2 = h * h;

    Image<T> result(H, W, ch, T(0));

    for (int y = 0; y < H; ++y) {
        for (int x = 0; x < W; ++x) {
            std::vector<double> sumC(ch, 0);
            double sumW = 0;

            // 探索窓
            int sy0 = std::max(0, y - sHalf), sy1 = std::min(H - 1, y + sHalf);
            int sx0 = std::max(0, x - sHalf), sx1 = std::min(W - 1, x + sHalf);

            for (int sy = sy0; sy <= sy1; ++sy) {
                for (int sx = sx0; sx <= sx1; ++sx) {
                    // パッチ間の二乗距離
                    double dist2 = 0;
                    int count = 0;
                    for (int ty = -tHalf; ty <= tHalf; ++ty) {
                        for (int tx = -tHalf; tx <= tHalf; ++tx) {
                            int py1 = y + ty, px1 = x + tx;
                            int py2 = sy + ty, px2 = sx + tx;
                            if (py1 < 0 || py1 >= H || px1 < 0 || px1 >= W) continue;
                            if (py2 < 0 || py2 >= H || px2 < 0 || px2 >= W) continue;
                            for (int c = 0; c < ch; ++c) {
                                double d = double(src.at(py1, px1, c)) - double(src.at(py2, px2, c));
                                dist2 += d * d;
                            }
                            ++count;
                        }
                    }
                    if (count > 0) dist2 /= count;

                    double w = std::exp(-dist2 / h2);
                    sumW += w;
                    for (int c = 0; c < ch; ++c)
                        sumC[c] += w * double(src.at(sy, sx, c));
                }
            }

            if (sumW > 0)
                for (int c = 0; c < ch; ++c)
                    result.at(y, x, c) = T(sumC[c] / sumW);
        }
    }
    return result;
}

} // namespace calx::image

#endif
