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

/**
 * @file distance_transform.hpp
 * @brief 距離変換 (OCV-7: distanceTransform)
 *
 * 各ピクセルから最寄りのゼロピクセルまでの距離を計算。
 * セグメンテーション、形態解析、Voronoi 分割の前処理。
 */

#ifndef CALX_IMAGE_IMGPROC_DISTANCE_TRANSFORM_HPP
#define CALX_IMAGE_IMGPROC_DISTANCE_TRANSFORM_HPP

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

namespace calx::image {

enum class DistanceType { L1, L2, Chessboard };

/// 距離変換 (2パス法)
/// 入力: 2値画像 (0=前景, 非0=背景)
/// 出力: 各ピクセルから最寄りの 0 ピクセルまでの距離
template<typename T>
Image<T> distanceTransform(const Image<T>& binary,
                            DistanceType type = DistanceType::L2) {
    int h = static_cast<int>(binary.height());
    int w = static_cast<int>(binary.width());
    constexpr T INF = T(1e9);

    Image<T> dist(h, w, 1, T(0));

    // 初期化: 0 ピクセル → 距離 0, それ以外 → INF
    for (int y = 0; y < h; ++y)
        for (int x = 0; x < w; ++x)
            dist(y, x) = (binary(y, x) == T(0)) ? T(0) : INF;

    if (type == DistanceType::L1 || type == DistanceType::Chessboard) {
        T d = (type == DistanceType::L1) ? T(1) : T(1);

        // Forward pass
        for (int y = 0; y < h; ++y)
            for (int x = 0; x < w; ++x) {
                if (y > 0) dist(y,x) = std::min(dist(y,x), dist(y-1,x) + T(1));
                if (x > 0) dist(y,x) = std::min(dist(y,x), dist(y,x-1) + T(1));
                if (type == DistanceType::Chessboard && y > 0 && x > 0)
                    dist(y,x) = std::min(dist(y,x), dist(y-1,x-1) + T(1));
                if (type == DistanceType::Chessboard && y > 0 && x < w-1)
                    dist(y,x) = std::min(dist(y,x), dist(y-1,x+1) + T(1));
            }

        // Backward pass
        for (int y = h-1; y >= 0; --y)
            for (int x = w-1; x >= 0; --x) {
                if (y < h-1) dist(y,x) = std::min(dist(y,x), dist(y+1,x) + T(1));
                if (x < w-1) dist(y,x) = std::min(dist(y,x), dist(y,x+1) + T(1));
                if (type == DistanceType::Chessboard && y < h-1 && x < w-1)
                    dist(y,x) = std::min(dist(y,x), dist(y+1,x+1) + T(1));
                if (type == DistanceType::Chessboard && y < h-1 && x > 0)
                    dist(y,x) = std::min(dist(y,x), dist(y+1,x-1) + T(1));
            }
    } else {
        // L2: 分離可能な二乗距離変換 (Meijster et al. 2000)
        // Phase 1: 各列の最寄り 0 行までの距離
        std::vector<T> col_dist(h);
        for (int x = 0; x < w; ++x) {
            // Forward
            col_dist[0] = (binary(0,x) == T(0)) ? T(0) : INF;
            for (int y = 1; y < h; ++y)
                col_dist[y] = (binary(y,x) == T(0)) ? T(0) : col_dist[y-1] + T(1);
            // Backward
            for (int y = h-2; y >= 0; --y)
                col_dist[y] = std::min(col_dist[y], col_dist[y+1] + T(1));
            for (int y = 0; y < h; ++y)
                dist(y,x) = col_dist[y] * col_dist[y];  // 二乗距離
        }

        // Phase 2: 各行の二乗距離変換
        std::vector<T> row_buf(w);
        for (int y = 0; y < h; ++y) {
            for (int x = 0; x < w; ++x) row_buf[x] = dist(y,x);
            for (int x = 0; x < w; ++x) {
                T best = row_buf[x];
                for (int dx = 1; dx < w; ++dx) {
                    T d2 = T(dx * dx);
                    if (d2 >= best) break;
                    if (x - dx >= 0) best = std::min(best, row_buf[x-dx] + d2);
                    if (x + dx < w)  best = std::min(best, row_buf[x+dx] + d2);
                }
                dist(y,x) = std::sqrt(best);
            }
        }
    }
    return dist;
}

} // namespace calx::image

#endif
