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

// histogram.hpp — ヒストグラム演算
//
// calcHist:      ヒストグラム計算 (256 bins for [0,1] range)
// equalizeHist:  ヒストグラム均一化
// clahe:         Contrast Limited Adaptive Histogram Equalization

#ifndef CALX_IMAGE_IMGPROC_HISTOGRAM_HPP
#define CALX_IMAGE_IMGPROC_HISTOGRAM_HPP

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

namespace calx::image {

    /// ヒストグラム計算
    /// @param img      入力画像
    /// @param channel  対象チャンネル
    /// @param numBins  ビン数 (デフォルト 256)
    /// @return         各ビンのピクセル数
    template<typename T>
    std::vector<size_t> calcHist(const Image<T>& img, int channel = 0, int numBins = 256) {
        assert(channel >= 0 && channel < img.numChannels());
        assert(numBins > 0);

        std::vector<size_t> hist(numBins, 0);
        size_t H = img.height(), W = img.width();
        const auto& ch = img.channel(channel);

        for (size_t y = 0; y < H; ++y) {
            for (size_t x = 0; x < W; ++x) {
                T val = ch(y, x);
                int bin = static_cast<int>(val * static_cast<T>(numBins));
                bin = std::clamp(bin, 0, numBins - 1);
                ++hist[bin];
            }
        }
        return hist;
    }

    /// ヒストグラム均一化
    /// @param img      入力画像
    /// @param channel  対象チャンネル
    /// @return         均一化された画像
    template<typename T>
    Image<T> equalizeHist(const Image<T>& img, int channel = 0) {
        assert(channel >= 0 && channel < img.numChannels());

        size_t H = img.height(), W = img.width();
        size_t totalPixels = H * W;
        if (totalPixels == 0) return img;

        const int numBins = 256;
        auto hist = calcHist(img, channel, numBins);

        // 累積分布関数 (CDF) を構築
        std::vector<size_t> cdf(numBins, 0);
        cdf[0] = hist[0];
        for (int i = 1; i < numBins; ++i)
            cdf[i] = cdf[i - 1] + hist[i];

        // CDF の最小非ゼロ値
        size_t cdfMin = 0;
        for (int i = 0; i < numBins; ++i) {
            if (cdf[i] > 0) { cdfMin = cdf[i]; break; }
        }

        // ルックアップテーブル
        std::vector<T> lut(numBins, T{0});
        if (totalPixels > cdfMin) {
            double denom = static_cast<double>(totalPixels - cdfMin);
            for (int i = 0; i < numBins; ++i) {
                if (cdf[i] > 0) {
                    lut[i] = static_cast<T>(
                        static_cast<double>(cdf[i] - cdfMin) / denom);
                }
            }
        }

        // 結果画像を構築
        Image<T> result = img;
        auto& resCh = result.channel(channel);
        for (size_t y = 0; y < H; ++y) {
            for (size_t x = 0; x < W; ++x) {
                int bin = static_cast<int>(img.channel(channel)(y, x) * static_cast<T>(numBins));
                bin = std::clamp(bin, 0, numBins - 1);
                resCh(y, x) = lut[bin];
            }
        }
        return result;
    }

    /// CLAHE (Contrast Limited Adaptive Histogram Equalization)
    /// @param img        入力画像
    /// @param tileSize   タイルのグリッド数 (画像を tileSize x tileSize に分割)
    /// @param clipLimit  クリッピング制限値 (1.0 = 制限なし)
    /// @param channel    対象チャンネル
    /// @return           CLAHE 適用後の画像
    template<typename T>
    Image<T> clahe(const Image<T>& img, int tileSize = 8, T clipLimit = T{2.0},
                   int channel = 0) {
        assert(channel >= 0 && channel < img.numChannels());
        assert(tileSize >= 1);

        size_t H = img.height(), W = img.width();
        if (H == 0 || W == 0) return img;

        const int numBins = 256;
        size_t tileH = (H + tileSize - 1) / tileSize;
        size_t tileW = (W + tileSize - 1) / tileSize;

        const auto& ch = img.channel(channel);

        // 各タイルのヒストグラムを計算し、クリッピングして CDF を作成
        // cdfTable[ty][tx] = CDF (numBins 個)
        std::vector<std::vector<std::vector<double>>> cdfTable(
            tileSize, std::vector<std::vector<double>>(tileSize, std::vector<double>(numBins, 0)));

        for (int ty = 0; ty < tileSize; ++ty) {
            for (int tx = 0; tx < tileSize; ++tx) {
                size_t y0 = ty * tileH;
                size_t x0 = tx * tileW;
                size_t y1 = std::min(y0 + tileH, H);
                size_t x1 = std::min(x0 + tileW, W);

                // タイルのヒストグラム
                std::vector<double> hist(numBins, 0);
                size_t count = 0;
                for (size_t y = y0; y < y1; ++y) {
                    for (size_t x = x0; x < x1; ++x) {
                        int bin = static_cast<int>(ch(y, x) * static_cast<T>(numBins));
                        bin = std::clamp(bin, 0, numBins - 1);
                        hist[bin] += 1.0;
                        ++count;
                    }
                }
                if (count == 0) continue;

                // クリッピング
                double clipVal = static_cast<double>(clipLimit)
                                 * static_cast<double>(count) / numBins;
                double excess = 0;
                for (int i = 0; i < numBins; ++i) {
                    if (hist[i] > clipVal) {
                        excess += hist[i] - clipVal;
                        hist[i] = clipVal;
                    }
                }
                // 余剰分を均等に再分配
                double incr = excess / numBins;
                for (int i = 0; i < numBins; ++i)
                    hist[i] += incr;

                // CDF を構築 (正規化)
                auto& cdf = cdfTable[ty][tx];
                cdf[0] = hist[0];
                for (int i = 1; i < numBins; ++i)
                    cdf[i] = cdf[i - 1] + hist[i];
                double total = cdf[numBins - 1];
                if (total > 0) {
                    for (int i = 0; i < numBins; ++i)
                        cdf[i] /= total;
                }
            }
        }

        // バイリニア補間でタイル境界を滑らかに接続
        Image<T> result = img;
        auto& resCh = result.channel(channel);

        for (size_t y = 0; y < H; ++y) {
            for (size_t x = 0; x < W; ++x) {
                int bin = static_cast<int>(ch(y, x) * static_cast<T>(numBins));
                bin = std::clamp(bin, 0, numBins - 1);

                // タイル中心座標での正規化位置
                double fy = (static_cast<double>(y) + 0.5) / static_cast<double>(tileH) - 0.5;
                double fx = (static_cast<double>(x) + 0.5) / static_cast<double>(tileW) - 0.5;

                int ty0 = std::clamp(static_cast<int>(std::floor(fy)), 0, tileSize - 1);
                int tx0 = std::clamp(static_cast<int>(std::floor(fx)), 0, tileSize - 1);
                int ty1 = std::clamp(ty0 + 1, 0, tileSize - 1);
                int tx1 = std::clamp(tx0 + 1, 0, tileSize - 1);

                double ay = fy - std::floor(fy);
                double ax = fx - std::floor(fx);

                // 4 つのタイル CDF から補間
                double v00 = cdfTable[ty0][tx0][bin];
                double v01 = cdfTable[ty0][tx1][bin];
                double v10 = cdfTable[ty1][tx0][bin];
                double v11 = cdfTable[ty1][tx1][bin];

                double val = (1 - ay) * ((1 - ax) * v00 + ax * v01)
                           +      ay  * ((1 - ax) * v10 + ax * v11);

                resCh(y, x) = static_cast<T>(std::clamp(val, 0.0, 1.0));
            }
        }
        return result;
    }

} // namespace calx::image

#endif // CALX_IMAGE_IMGPROC_HISTOGRAM_HPP
