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

/**
 * @file watershed.hpp
 * @brief Watershed / GrabCut セグメンテーション (OCV-11)
 */

#ifndef CALX_IMAGE_IMGPROC_WATERSHED_HPP
#define CALX_IMAGE_IMGPROC_WATERSHED_HPP

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

namespace calx::image {

/// Watershed セグメンテーション (マーカーベース)
/// markers: 初期マーカー画像 (0=未定, >0=領域ID, -1=境界)
/// 出力: 更新された markers (-1=境界線)
template<typename T>
Image<int> watershed(const Image<T>& img, Image<int>& markers) {
    int h = static_cast<int>(img.height());
    int w = static_cast<int>(img.width());
    int ch = std::min(img.numChannels(), 3);
    constexpr int BOUNDARY = -1;

    // 勾配画像 (エッジ強度)
    auto gradient = [&](int y, int x) -> double {
        double g = 0;
        for (auto [dy,dx] : std::vector<std::pair<int,int>>{{0,1},{1,0}}) {
            int ny = y+dy, nx = x+dx;
            if (ny >= h || nx >= w) continue;
            for (int c = 0; c < ch; ++c) {
                double d = double(img.at(y,x,c)) - double(img.at(ny,nx,c));
                g += d * d;
            }
        }
        return std::sqrt(g);
    };

    // 優先度キュー (勾配が小さい順に flood)
    using PQE = std::tuple<double, int, int>;
    std::priority_queue<PQE, std::vector<PQE>, std::greater<>> pq;

    // マーカー付きピクセルの隣接をキューに追加
    for (int y = 0; y < h; ++y)
        for (int x = 0; x < w; ++x)
            if (markers(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) continue;
                    if (markers(ny, nx) == 0)
                        pq.push({gradient(ny, nx), ny, nx});
                }
            }

    // Flooding
    while (!pq.empty()) {
        auto [g, cy, cx] = pq.top(); pq.pop();
        if (markers(cy, cx) != 0) continue;

        // 隣接マーカーを調べる
        int label = 0;
        bool boundary = false;
        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;
            int m = markers(ny, nx);
            if (m > 0) {
                if (label == 0) label = m;
                else if (label != m) { boundary = true; break; }
            }
        }

        if (boundary) {
            markers(cy, cx) = BOUNDARY;
        } else if (label > 0) {
            markers(cy, cx) = label;
            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 (markers(ny, nx) == 0)
                    pq.push({gradient(ny, nx), ny, nx});
            }
        }
    }
    return markers;
}

/// GrabCut セグメンテーション (簡易版: 反復的前景/背景分離)
/// mask: 0=BG_確定, 1=FG_確定, 2=BG_候補, 3=FG_候補
/// rect: 初期前景矩形 (mask の 3 に設定)
template<typename T>
void grabCut(const Image<T>& img, Image<int>& mask,
              int rx, int ry, int rw, int rh, int iterations = 5) {
    int h = static_cast<int>(img.height());
    int w = static_cast<int>(img.width());
    int ch = std::min(img.numChannels(), 3);

    // 初期マスク設定
    for (int y = 0; y < h; ++y)
        for (int x = 0; x < w; ++x) {
            if (y >= ry && y < ry+rh && x >= rx && x < rx+rw)
                mask(y, x) = 3;  // 前景候補
            else
                mask(y, x) = 0;  // 背景確定
        }

    // 反復: 前景/背景の平均色を計算し、距離でマスクを更新
    for (int iter = 0; iter < iterations; ++iter) {
        // 前景・背景の平均色
        std::vector<double> fgMean(ch, 0), bgMean(ch, 0);
        int fgCount = 0, bgCount = 0;
        for (int y = 0; y < h; ++y)
            for (int x = 0; x < w; ++x) {
                int m = mask(y, x);
                if (m == 1 || m == 3) {
                    for (int c = 0; c < ch; ++c) fgMean[c] += double(img.at(y,x,c));
                    ++fgCount;
                } else {
                    for (int c = 0; c < ch; ++c) bgMean[c] += double(img.at(y,x,c));
                    ++bgCount;
                }
            }
        if (fgCount > 0) for (auto& v : fgMean) v /= fgCount;
        if (bgCount > 0) for (auto& v : bgMean) v /= bgCount;

        // 候補ピクセルの再分類
        for (int y = 0; y < h; ++y)
            for (int x = 0; x < w; ++x) {
                int m = mask(y, x);
                if (m == 0 || m == 1) continue;  // 確定済みはスキップ
                double dFg = 0, dBg = 0;
                for (int c = 0; c < ch; ++c) {
                    double vf = double(img.at(y,x,c)) - fgMean[c];
                    double vb = double(img.at(y,x,c)) - bgMean[c];
                    dFg += vf * vf;
                    dBg += vb * vb;
                }
                mask(y, x) = (dFg < dBg) ? 3 : 2;
            }
    }
}

} // namespace calx::image

#endif
