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

// metrics.hpp — 画像品質指標
//
// PSNR: Peak Signal-to-Noise Ratio
// SSIM: Structural Similarity Index

#ifndef CALX_IMAGE_QUALITY_METRICS_HPP
#define CALX_IMAGE_QUALITY_METRICS_HPP

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

namespace calx::image {

    /// PSNR (Peak Signal-to-Noise Ratio)
    /// @param a       画像 A
    /// @param b       画像 B (同サイズ)
    /// @param maxVal  ピクセル最大値 (デフォルト 1.0)
    /// @return        PSNR (dB), 同一画像なら +inf
    template<typename T>
    T psnr(const Image<T>& a, const Image<T>& b, T maxVal = T{1.0}) {
        assert(a.height() == b.height() && a.width() == b.width());
        assert(a.numChannels() == b.numChannels());

        size_t H = a.height(), W = a.width();
        int nch = a.numChannels();
        double mse = 0;
        size_t count = 0;

        for (int c = 0; c < nch; ++c) {
            const auto& aCh = a.channel(c);
            const auto& bCh = b.channel(c);
            for (size_t y = 0; y < H; ++y) {
                for (size_t x = 0; x < W; ++x) {
                    double diff = static_cast<double>(aCh(y, x)) - static_cast<double>(bCh(y, x));
                    mse += diff * diff;
                    ++count;
                }
            }
        }

        if (count == 0) return T{0};
        mse /= count;

        if (mse < 1e-15)
            return std::numeric_limits<T>::infinity();

        double maxD = static_cast<double>(maxVal);
        return static_cast<T>(10.0 * std::log10(maxD * maxD / mse));
    }

    /// SSIM (Structural Similarity Index)
    /// Wang et al., "Image Quality Assessment: From Error Visibility to
    /// Structural Similarity", IEEE TIP 2004.
    /// @param a           画像 A (1ch)
    /// @param b           画像 B (1ch, 同サイズ)
    /// @param windowSize  ガウシアンウィンドウサイズ (奇数)
    /// @return            SSIM 値 [-1, 1] (1 = 完全一致)
    template<typename T>
    T ssim(const Image<T>& a, const Image<T>& b, int windowSize = 11) {
        assert(a.height() == b.height() && a.width() == b.width());
        assert(a.numChannels() >= 1 && b.numChannels() >= 1);
        assert(windowSize > 0 && windowSize % 2 == 1);

        int H = static_cast<int>(a.height());
        int W = static_cast<int>(b.width());
        int half = windowSize / 2;

        const auto& aCh = a.channel(0);
        const auto& bCh = b.channel(0);

        // 安定化定数 (L=1.0 (maxVal), K1=0.01, K2=0.03)
        constexpr double K1 = 0.01, K2 = 0.03;
        constexpr double L = 1.0;
        double C1 = (K1 * L) * (K1 * L);
        double C2 = (K2 * L) * (K2 * L);

        // ガウシアンカーネル生成
        double sigma = static_cast<double>(windowSize) / 6.0;
        std::vector<double> kernel(windowSize);
        double ksum = 0;
        for (int i = 0; i < windowSize; ++i) {
            double x = i - half;
            kernel[i] = std::exp(-x * x / (2 * sigma * sigma));
            ksum += kernel[i];
        }
        for (auto& k : kernel) k /= ksum;

        // 2D ガウシアン重み
        std::vector<std::vector<double>> weight(windowSize, std::vector<double>(windowSize));
        for (int i = 0; i < windowSize; ++i)
            for (int j = 0; j < windowSize; ++j)
                weight[i][j] = kernel[i] * kernel[j];

        double totalSSIM = 0;
        int count = 0;

        for (int y = half; y < H - half; ++y) {
            for (int x = half; x < W - half; ++x) {
                double muA = 0, muB = 0;
                double sigA2 = 0, sigB2 = 0, sigAB = 0;

                // 局所統計量の計算
                for (int ky = -half; ky <= half; ++ky) {
                    for (int kx = -half; kx <= half; ++kx) {
                        double w = weight[ky + half][kx + half];
                        double va = static_cast<double>(aCh(y + ky, x + kx));
                        double vb = static_cast<double>(bCh(y + ky, x + kx));
                        muA += w * va;
                        muB += w * vb;
                    }
                }

                for (int ky = -half; ky <= half; ++ky) {
                    for (int kx = -half; kx <= half; ++kx) {
                        double w = weight[ky + half][kx + half];
                        double da = static_cast<double>(aCh(y + ky, x + kx)) - muA;
                        double db = static_cast<double>(bCh(y + ky, x + kx)) - muB;
                        sigA2 += w * da * da;
                        sigB2 += w * db * db;
                        sigAB += w * da * db;
                    }
                }

                double num = (2 * muA * muB + C1) * (2 * sigAB + C2);
                double den = (muA * muA + muB * muB + C1) * (sigA2 + sigB2 + C2);
                totalSSIM += num / den;
                ++count;
            }
        }

        return (count > 0) ? static_cast<T>(totalSSIM / count) : T{1};
    }

} // namespace calx::image

#endif // CALX_IMAGE_QUALITY_METRICS_HPP
