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

// ndimage.hpp — N 次元画像フィルタ (scipy.ndimage 対応)
//
// 新規実装:
//   - convolve (畳み込み), gaussianFilter, medianFilter
//   - Sobel/Prewitt エッジ検出, ラプラシアンフィルタ
//   - calx Matrix<T> ベース

#ifndef CALX_IMAGE_FILTERS_NDIMAGE_HPP
#define CALX_IMAGE_FILTERS_NDIMAGE_HPP

#include <math/core/matrix.hpp>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <vector>

namespace calx::image {

    /// 2D 畳み込み (zero-padding)
    /// @param image   入力画像 (H × W)
    /// @param kernel  カーネル (kH × kW, 奇数サイズ)
    /// @return        畳み込み結果 (H × W)
    template<typename T>
    Matrix<T> convolve(const Matrix<T>& image, const Matrix<T>& kernel) {
        int H = static_cast<int>(image.rows());
        int W = static_cast<int>(image.cols());
        int kH = static_cast<int>(kernel.rows());
        int kW = static_cast<int>(kernel.cols());
        int hkH = kH / 2;
        int hkW = kW / 2;

        Matrix<T> result(H, W, T{0});

        for (int y = 0; y < H; ++y) {
            for (int x = 0; x < W; ++x) {
                T sum = T{0};
                for (int ky = 0; ky < kH; ++ky) {
                    for (int kx = 0; kx < kW; ++kx) {
                        int iy = y + ky - hkH;
                        int ix = x + kx - hkW;
                        if (iy >= 0 && iy < H && ix >= 0 && ix < W) {
                            sum += image(iy, ix) * kernel(ky, kx);
                        }
                    }
                }
                result(y, x) = sum;
            }
        }
        return result;
    }

    /// ガウシアンカーネル生成
    template<typename T>
    Matrix<T> gaussianKernel(int size, T sigma) {
        assert(size > 0 && size % 2 == 1);
        int half = size / 2;
        Matrix<T> kernel(size, size);
        T sum = T{0};

        for (int y = -half; y <= half; ++y) {
            for (int x = -half; x <= half; ++x) {
                T val = std::exp(-(x * x + y * y) / (T{2} * sigma * sigma));
                kernel(y + half, x + half) = val;
                sum += val;
            }
        }

        // 正規化
        for (int y = 0; y < size; ++y)
            for (int x = 0; x < size; ++x)
                kernel(y, x) /= sum;

        return kernel;
    }

    /// ガウシアンフィルタ
    template<typename T>
    Matrix<T> gaussianFilter(const Matrix<T>& image, T sigma, int size = 0) {
        if (size == 0) size = static_cast<int>(std::ceil(sigma * 6)) | 1;  // 奇数に
        return convolve(image, gaussianKernel<T>(size, sigma));
    }

    /// メディアンフィルタ
    template<typename T>
    Matrix<T> medianFilter(const Matrix<T>& image, int size = 3) {
        assert(size > 0 && size % 2 == 1);
        int H = static_cast<int>(image.rows());
        int W = static_cast<int>(image.cols());
        int half = size / 2;

        Matrix<T> result(H, W);
        std::vector<T> window;

        for (int y = 0; y < H; ++y) {
            for (int x = 0; x < W; ++x) {
                window.clear();
                for (int ky = -half; ky <= half; ++ky) {
                    for (int kx = -half; kx <= half; ++kx) {
                        int iy = y + ky;
                        int ix = x + kx;
                        if (iy >= 0 && iy < H && ix >= 0 && ix < W) {
                            window.push_back(image(iy, ix));
                        }
                    }
                }
                std::nth_element(window.begin(), window.begin() + window.size() / 2, window.end());
                result(y, x) = window[window.size() / 2];
            }
        }
        return result;
    }

    /// Sobel フィルタ (水平方向微分)
    template<typename T>
    Matrix<T> sobelX() {
        Matrix<T> k(3, 3);
        k(0,0) = T{-1}; k(0,1) = T{0}; k(0,2) = T{1};
        k(1,0) = T{-2}; k(1,1) = T{0}; k(1,2) = T{2};
        k(2,0) = T{-1}; k(2,1) = T{0}; k(2,2) = T{1};
        return k;
    }

    /// Sobel フィルタ (垂直方向微分)
    template<typename T>
    Matrix<T> sobelY() {
        Matrix<T> k(3, 3);
        k(0,0) = T{-1}; k(0,1) = T{-2}; k(0,2) = T{-1};
        k(1,0) = T{0};  k(1,1) = T{0};  k(1,2) = T{0};
        k(2,0) = T{1};  k(2,1) = T{2};  k(2,2) = T{1};
        return k;
    }

    /// Sobel エッジ検出 (勾配の大きさを返す)
    template<typename T>
    Matrix<T> sobelEdge(const Matrix<T>& image) {
        auto gx = convolve(image, sobelX<T>());
        auto gy = convolve(image, sobelY<T>());

        int H = static_cast<int>(image.rows());
        int W = static_cast<int>(image.cols());
        Matrix<T> result(H, W);

        for (int y = 0; y < H; ++y)
            for (int x = 0; x < W; ++x)
                result(y, x) = std::sqrt(gx(y, x) * gx(y, x) + gy(y, x) * gy(y, x));

        return result;
    }

    /// Prewitt フィルタ (水平)
    template<typename T>
    Matrix<T> prewittX() {
        Matrix<T> k(3, 3);
        k(0,0) = T{-1}; k(0,1) = T{0}; k(0,2) = T{1};
        k(1,0) = T{-1}; k(1,1) = T{0}; k(1,2) = T{1};
        k(2,0) = T{-1}; k(2,1) = T{0}; k(2,2) = T{1};
        return k;
    }

    /// Prewitt フィルタ (垂直)
    template<typename T>
    Matrix<T> prewittY() {
        Matrix<T> k(3, 3);
        k(0,0) = T{-1}; k(0,1) = T{-1}; k(0,2) = T{-1};
        k(1,0) = T{0};  k(1,1) = T{0};  k(1,2) = T{0};
        k(2,0) = T{1};  k(2,1) = T{1};  k(2,2) = T{1};
        return k;
    }

    /// Prewitt エッジ検出
    template<typename T>
    Matrix<T> prewittEdge(const Matrix<T>& image) {
        auto gx = convolve(image, prewittX<T>());
        auto gy = convolve(image, prewittY<T>());

        int H = static_cast<int>(image.rows());
        int W = static_cast<int>(image.cols());
        Matrix<T> result(H, W);

        for (int y = 0; y < H; ++y)
            for (int x = 0; x < W; ++x)
                result(y, x) = std::sqrt(gx(y, x) * gx(y, x) + gy(y, x) * gy(y, x));

        return result;
    }

    /// ラプラシアンカーネル
    template<typename T>
    Matrix<T> laplacianKernel() {
        Matrix<T> k(3, 3, T{0});
        k(0,1) = T{1};
        k(1,0) = T{1}; k(1,1) = T{-4}; k(1,2) = T{1};
        k(2,1) = T{1};
        return k;
    }

    /// ラプラシアンフィルタ
    template<typename T>
    Matrix<T> laplacianFilter(const Matrix<T>& image) {
        return convolve(image, laplacianKernel<T>());
    }

} // namespace calx::image

#endif // CALX_IMAGE_FILTERS_NDIMAGE_HPP
