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

/**
 * @file superpixels.hpp
 * @brief SLIC 超画素 (OCV-33: Superpixels)
 *
 * Simple Linear Iterative Clustering (Achanta et al. 2012)
 */

#ifndef CALX_IMAGE_IMGPROC_SUPERPIXELS_HPP
#define CALX_IMAGE_IMGPROC_SUPERPIXELS_HPP

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

namespace calx::image {

/// SLIC 超画素分割
/// numSuperpixels: 目標超画素数, compactness: 空間 vs 色の重み (10-40)
template<typename T>
Image<int> slic(const Image<T>& img, int numSuperpixels = 200,
                 double compactness = 10.0, int iterations = 10) {
    int h = static_cast<int>(img.height());
    int w = static_cast<int>(img.width());
    int ch = std::min(img.numChannels(), 3);
    int S = static_cast<int>(std::sqrt(double(h * w) / numSuperpixels));
    if (S < 1) S = 1;

    // クラスタ中心を初期化 (グリッド上)
    struct Center { double x, y; std::vector<double> color; };
    std::vector<Center> centers;
    for (int cy = S / 2; cy < h; cy += S)
        for (int cx = S / 2; cx < w; cx += S) {
            Center c;
            c.x = cx; c.y = cy;
            c.color.resize(ch);
            for (int k = 0; k < ch; ++k)
                c.color[k] = double(img.at(cy, cx, k));
            centers.push_back(c);
        }

    int K = static_cast<int>(centers.size());
    std::vector<int> labels(h * w, -1);
    std::vector<double> dist(h * w, std::numeric_limits<double>::max());

    // 反復
    for (int iter = 0; iter < iterations; ++iter) {
        // 各中心の 2S × 2S 近傍を割り当て
        std::fill(dist.begin(), dist.end(), std::numeric_limits<double>::max());

        for (int k = 0; k < K; ++k) {
            int cy = static_cast<int>(centers[k].y);
            int cx = static_cast<int>(centers[k].x);
            int y0 = std::max(0, cy - S), y1 = std::min(h - 1, cy + S);
            int x0 = std::max(0, cx - S), x1 = std::min(w - 1, cx + S);

            for (int y = y0; y <= y1; ++y)
                for (int x = x0; x <= x1; ++x) {
                    // 色距離
                    double dc = 0;
                    for (int c = 0; c < ch; ++c) {
                        double d = double(img.at(y, x, c)) - centers[k].color[c];
                        dc += d * d;
                    }
                    dc = std::sqrt(dc);

                    // 空間距離
                    double ds = std::sqrt(double((x - cx)*(x - cx) + (y - cy)*(y - cy)));

                    double D = dc + (compactness / S) * ds;

                    int idx = y * w + x;
                    if (D < dist[idx]) {
                        dist[idx] = D;
                        labels[idx] = k;
                    }
                }
        }

        // 中心を更新
        std::vector<double> sumX(K, 0), sumY(K, 0);
        std::vector<std::vector<double>> sumC(K, std::vector<double>(ch, 0));
        std::vector<int> count(K, 0);

        for (int y = 0; y < h; ++y)
            for (int x = 0; x < w; ++x) {
                int k = labels[y * w + x];
                if (k < 0) continue;
                sumX[k] += x; sumY[k] += y;
                for (int c = 0; c < ch; ++c)
                    sumC[k][c] += double(img.at(y, x, c));
                ++count[k];
            }

        for (int k = 0; k < K; ++k) {
            if (count[k] > 0) {
                centers[k].x = sumX[k] / count[k];
                centers[k].y = sumY[k] / count[k];
                for (int c = 0; c < ch; ++c)
                    centers[k].color[c] = sumC[k][c] / count[k];
            }
        }
    }

    // ラベルマップを Image<int> として返す
    Image<int> labelMap(h, w);
    for (int y = 0; y < h; ++y)
        for (int x = 0; x < w; ++x)
            labelMap(y, x) = labels[y * w + x];
    return labelMap;
}

} // namespace calx::image

#endif
