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

/**
 * @file inpainting.hpp
 * @brief 画像修復 (OCV-23: Inpainting)
 *
 * Navier-Stokes (NS) ベースと Telea (FMM) ベースの 2 手法。
 */

#ifndef CALX_IMAGE_IMGPROC_INPAINTING_HPP
#define CALX_IMAGE_IMGPROC_INPAINTING_HPP

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

namespace calx::image {

enum class InpaintMethod { Telea, NS };

/// 画像修復 (Telea FMM ベース)
/// mask: 非0のピクセルを修復対象とする
template<typename T>
Image<T> inpaint(const Image<T>& src, const Image<T>& mask, double radius,
                  InpaintMethod method = InpaintMethod::Telea) {
    int h = static_cast<int>(src.height());
    int w = static_cast<int>(src.width());
    int ch = src.numChannels();
    Image<T> result = src;

    // 既知/未知のマップ
    std::vector<std::vector<bool>> known(h, std::vector<bool>(w, true));
    for (int y = 0; y < h; ++y)
        for (int x = 0; x < w; ++x)
            if (mask(y, x) != T(0)) known[y][x] = false;

    // 距離マップ (FMM 用)
    constexpr double INF = 1e9;
    std::vector<std::vector<double>> dist(h, std::vector<double>(w, INF));

    // 境界ピクセルをキューに追加
    using PQEntry = std::tuple<double, int, int>;
    std::priority_queue<PQEntry, std::vector<PQEntry>, std::greater<>> pq;

    for (int y = 0; y < h; ++y)
        for (int x = 0; x < w; ++x)
            if (known[y][x]) {
                dist[y][x] = 0;
                // 隣接に未知があれば境界
                for (auto [dy,dx] : std::vector<std::pair<int,int>>{{-1,0},{1,0},{0,-1},{0,1}}) {
                    int ny = y+dy, nx = x+dx;
                    if (ny >= 0 && ny < h && nx >= 0 && nx < w && !known[ny][nx])
                        pq.push({0, y, x});
                }
            }

    // Fast Marching Method で外側から内側へ修復
    while (!pq.empty()) {
        auto [d, cy, cx] = pq.top(); pq.pop();
        if (d > dist[cy][cx] + 1e-6) continue;

        for (auto [dy,dx] : std::vector<std::pair<int,int>>{{-1,0},{1,0},{0,-1},{0,1}}) {
            int ny = cy+dy, nx = cx+dx;
            if (ny < 0 || ny >= h || nx < 0 || nx >= w) continue;
            if (known[ny][nx]) continue;

            double newDist = dist[cy][cx] + 1.0;
            if (newDist >= dist[ny][nx]) continue;
            dist[ny][nx] = newDist;

            // 修復: radius 内の既知ピクセルの重み付き平均
            double sumW = 0;
            std::vector<double> sumC(ch, 0);
            int r = static_cast<int>(radius);
            for (int ky = std::max(0, ny-r); ky <= std::min(h-1, ny+r); ++ky)
                for (int kx = std::max(0, nx-r); kx <= std::min(w-1, nx+r); ++kx) {
                    if (!known[ky][kx]) continue;
                    double dd = std::sqrt(double((ky-ny)*(ky-ny) + (kx-nx)*(kx-nx)));
                    if (dd > radius) continue;
                    double weight = 1.0 / (dd + 1e-6);
                    sumW += weight;
                    for (int c = 0; c < ch; ++c)
                        sumC[c] += weight * double(result.at(ky, kx, c));
                }

            if (sumW > 0) {
                for (int c = 0; c < ch; ++c)
                    result.at(ny, nx, c) = T(sumC[c] / sumW);
            }
            known[ny][nx] = true;
            pq.push({newDist, ny, nx});
        }
    }
    return result;
}

} // namespace calx::image

#endif
