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

// stereo_matching.hpp — ステレオマッチング (BM / SGBM)
//
// 新規実装:
//   - ブロックマッチング (SAD ベース)
//   - Semi-Global Block Matching (4/8 パス DP コスト集約)
//   - 視差マップ推定、左右整合性チェック

#ifndef CALX_IMAGE_CV_STEREO_MATCHING_HPP
#define CALX_IMAGE_CV_STEREO_MATCHING_HPP

#include <image/Image.hpp>
#include <math/core/matrix.hpp>
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <limits>
#include <vector>

namespace calx::cv {

    /// ステレオマッチング結果
    template<typename T>
    struct StereoResult {
        calx::Matrix<T> disparity;  // 視差マップ (H x W)
    };

    namespace detail_stereo {

        /// SAD (Sum of Absolute Differences) を計算
        template<typename T>
        double computeSAD(const calx::Matrix<T>& left, const calx::Matrix<T>& right,
                          int ly, int lx, int ry, int rx, int halfBlock) {
            int H = static_cast<int>(left.rows());
            int W = static_cast<int>(left.cols());
            int rW = static_cast<int>(right.cols());
            double sad = 0;
            for (int dy = -halfBlock; dy <= halfBlock; ++dy) {
                for (int dx = -halfBlock; dx <= halfBlock; ++dx) {
                    int lyy = ly + dy, lxx = lx + dx;
                    int ryy = ry + dy, rxx = rx + dx;
                    if (lyy < 0 || lyy >= H || lxx < 0 || lxx >= W) continue;
                    if (ryy < 0 || ryy >= H || rxx < 0 || rxx >= rW) continue;
                    sad += std::abs(static_cast<double>(left(lyy, lxx)) -
                                   static_cast<double>(right(ryy, rxx)));
                }
            }
            return sad;
        }

        /// 3 点放物線サブピクセル補間
        template<typename T>
        T subpixelRefine(double costLeft, double costCenter, double costRight) {
            double denom = costLeft - 2.0 * costCenter + costRight;
            if (std::abs(denom) < 1e-12) return T{0};
            return static_cast<T>((costLeft - costRight) / (2.0 * denom));
        }

    } // namespace detail_stereo

    /// ブロックマッチング (BM)
    ///
    /// 各ピクセルについてエピポーラ線 (水平) 上でブロックの SAD を比較し、
    /// 最も一致する位置の視差を返す。
    ///
    /// @param left         左画像 (1ch グレースケール)
    /// @param right        右画像 (1ch グレースケール)
    /// @param blockSize    マッチングブロックのサイズ (奇数)
    /// @param maxDisparity 最大視差 (探索範囲)
    /// @return             視差マップ
    template<typename T>
    StereoResult<T> stereoBlockMatch(const calx::image::Image<T>& left,
                                      const calx::image::Image<T>& right,
                                      int blockSize = 9,
                                      int maxDisparity = 64) {
        if (blockSize % 2 == 0) ++blockSize;
        int halfBlock = blockSize / 2;

        const auto& lImg = left.channel(0);
        const auto& rImg = right.channel(0);
        int H = static_cast<int>(lImg.rows());
        int W = static_cast<int>(lImg.cols());

        calx::Matrix<T> disparity(H, W, T{0});

        for (int y = halfBlock; y < H - halfBlock; ++y) {
            for (int x = halfBlock; x < W - halfBlock; ++x) {
                double bestCost = std::numeric_limits<double>::max();
                int bestD = 0;
                double costPrev = 0, costBest = 0, costNext = 0;

                int dMax = std::min(maxDisparity, x - halfBlock);
                for (int d = 0; d <= dMax; ++d) {
                    double cost = detail_stereo::computeSAD(lImg, rImg, y, x, y, x - d, halfBlock);
                    if (cost < bestCost) {
                        bestCost = cost;
                        bestD = d;
                    }
                }

                // サブピクセル補間
                T subpixel = T{0};
                if (bestD > 0 && bestD < dMax) {
                    double cL = detail_stereo::computeSAD(lImg, rImg, y, x, y, x - (bestD - 1), halfBlock);
                    double cR = detail_stereo::computeSAD(lImg, rImg, y, x, y, x - (bestD + 1), halfBlock);
                    subpixel = detail_stereo::subpixelRefine<T>(cL, bestCost, cR);
                }

                disparity(y, x) = static_cast<T>(bestD) + subpixel;
            }
        }

        return {std::move(disparity)};
    }

    /// Semi-Global Block Matching (SGBM)
    ///
    /// 8 方向のパスで DP コストを集約し、ブロックマッチングより
    /// 滑らかで精度の高い視差マップを生成する。
    ///
    /// @param left         左画像 (1ch グレースケール)
    /// @param right        右画像 (1ch グレースケール)
    /// @param blockSize    マッチングブロックのサイズ (奇数)
    /// @param maxDisparity 最大視差 (探索範囲)
    /// @param P1           隣接ピクセル間の小さな視差変化ペナルティ
    /// @param P2           隣接ピクセル間の大きな視差変化ペナルティ
    /// @return             視差マップ
    template<typename T>
    StereoResult<T> stereoSGBM(const calx::image::Image<T>& left,
                                const calx::image::Image<T>& right,
                                int blockSize = 5,
                                int maxDisparity = 64,
                                int P1 = 8, int P2 = 32) {
        if (blockSize % 2 == 0) ++blockSize;
        int halfBlock = blockSize / 2;

        const auto& lImg = left.channel(0);
        const auto& rImg = right.channel(0);
        int H = static_cast<int>(lImg.rows());
        int W = static_cast<int>(lImg.cols());
        int D = maxDisparity;

        // --- Step 1: ピクセルごとのマッチングコスト (SAD) ---
        // cost[y][x][d] — 3D テーブルを 1D ベクトルで管理
        std::vector<double> pixelCost(static_cast<std::size_t>(H) * W * D,
                                      std::numeric_limits<double>::max());

        auto costIdx = [&](int y, int x, int d) -> std::size_t {
            return (static_cast<std::size_t>(y) * W + x) * D + d;
        };

        for (int y = halfBlock; y < H - halfBlock; ++y) {
            for (int x = halfBlock; x < W - halfBlock; ++x) {
                int dMax = std::min(D, x - halfBlock);
                for (int d = 0; d < dMax; ++d) {
                    pixelCost[costIdx(y, x, d)] =
                        detail_stereo::computeSAD(lImg, rImg, y, x, y, x - d, halfBlock);
                }
            }
        }

        // --- Step 2: 8 方向パスで DP コスト集約 ---
        // 方向ベクトル: 右、左、下、上、右下、左上、右上、左下
        static constexpr int dirs[][2] = {
            {0, 1}, {0, -1}, {1, 0}, {-1, 0},
            {1, 1}, {-1, -1}, {1, -1}, {-1, 1}
        };
        constexpr int numDirs = 8;

        // 集約コスト (全パスの合計)
        std::vector<double> aggregated(static_cast<std::size_t>(H) * W * D, 0.0);

        for (int dir = 0; dir < numDirs; ++dir) {
            int dy = dirs[dir][0];
            int dx = dirs[dir][1];

            // 前のピクセルの Lr 値を保持するバッファ
            // パスの走査順序を決定
            int yStart = (dy > 0) ? 0 : (dy < 0) ? H - 1 : 0;
            int yEnd   = (dy > 0) ? H : (dy < 0) ? -1    : H;
            int yStep  = (dy >= 0) ? 1 : -1;

            int xStart = (dx > 0) ? 0 : (dx < 0) ? W - 1 : 0;
            int xEnd   = (dx > 0) ? W : (dx < 0) ? -1    : W;
            int xStep  = (dx >= 0) ? 1 : -1;

            // 1 パス分のバッファ
            std::vector<double> Lr(static_cast<std::size_t>(H) * W * D, 0.0);

            for (int y = yStart; y != yEnd; y += yStep) {
                for (int x = xStart; x != xEnd; x += xStep) {
                    int py = y - dy;
                    int px = x - dx;
                    bool hasPrev = (py >= 0 && py < H && px >= 0 && px < W);

                    double prevMin = std::numeric_limits<double>::max();
                    if (hasPrev) {
                        for (int d = 0; d < D; ++d) {
                            double v = Lr[costIdx(py, px, d)];
                            if (v < prevMin) prevMin = v;
                        }
                    }

                    for (int d = 0; d < D; ++d) {
                        double C = pixelCost[costIdx(y, x, d)];
                        if (C >= std::numeric_limits<double>::max() * 0.5) {
                            Lr[costIdx(y, x, d)] = C;
                            continue;
                        }

                        if (!hasPrev) {
                            Lr[costIdx(y, x, d)] = C;
                        } else {
                            // DP 遷移:
                            // Lr(d) = C(d) + min(Lr_prev(d),
                            //                    Lr_prev(d-1)+P1,
                            //                    Lr_prev(d+1)+P1,
                            //                    min_k Lr_prev(k)+P2) - min_k Lr_prev(k)
                            double best = prevMin + P2;

                            double same = Lr[costIdx(py, px, d)];
                            if (same < best) best = same;

                            if (d > 0) {
                                double dm1 = Lr[costIdx(py, px, d - 1)] + P1;
                                if (dm1 < best) best = dm1;
                            }
                            if (d < D - 1) {
                                double dp1 = Lr[costIdx(py, px, d + 1)] + P1;
                                if (dp1 < best) best = dp1;
                            }

                            Lr[costIdx(y, x, d)] = C + best - prevMin;
                        }
                    }
                }
            }

            // 集約に加算
            for (std::size_t i = 0; i < aggregated.size(); ++i) {
                aggregated[i] += Lr[i];
            }
        }

        // --- Step 3: Winner-Takes-All で視差を選択 ---
        calx::Matrix<T> disparity(H, W, T{0});
        for (int y = 0; y < H; ++y) {
            for (int x = 0; x < W; ++x) {
                double bestCost = std::numeric_limits<double>::max();
                int bestD = 0;
                for (int d = 0; d < D; ++d) {
                    double c = aggregated[costIdx(y, x, d)];
                    if (c < bestCost) {
                        bestCost = c;
                        bestD = d;
                    }
                }

                // サブピクセル補間
                T subpixel = T{0};
                if (bestD > 0 && bestD < D - 1) {
                    double cL = aggregated[costIdx(y, x, bestD - 1)];
                    double cR = aggregated[costIdx(y, x, bestD + 1)];
                    subpixel = detail_stereo::subpixelRefine<T>(cL, bestCost, cR);
                }

                disparity(y, x) = static_cast<T>(bestD) + subpixel;
            }
        }

        // --- Step 4: 左右整合性チェック (オプション) ---
        // 右画像基準で逆方向マッチングし、不一致ピクセルを無効化
        calx::Matrix<T> disparityR(H, W, T{0});
        for (int y = 0; y < H; ++y) {
            for (int x = 0; x < W; ++x) {
                double bestCost = std::numeric_limits<double>::max();
                int bestD = 0;
                for (int d = 0; d < D; ++d) {
                    int lx = x + d;
                    if (lx >= W) break;
                    double c = aggregated[costIdx(y, lx, d)];
                    if (c < bestCost) {
                        bestCost = c;
                        bestD = d;
                    }
                }
                disparityR(y, x) = static_cast<T>(bestD);
            }
        }

        // 左右チェック: |dL(x) - dR(x - dL(x))| > 1 のピクセルを無効 (0) に
        for (int y = 0; y < H; ++y) {
            for (int x = 0; x < W; ++x) {
                T dL = disparity(y, x);
                int rx = x - static_cast<int>(dL);
                if (rx < 0 || rx >= W) {
                    disparity(y, x) = T{0};
                    continue;
                }
                T dR = disparityR(y, rx);
                if (std::abs(static_cast<double>(dL) - static_cast<double>(dR)) > 1.0) {
                    disparity(y, x) = T{0};
                }
            }
        }

        return {std::move(disparity)};
    }

    /// 視差マップをカラーで可視化
    template<typename T>
    calx::image::Image<T> drawDisparityMap(const calx::Matrix<T>& disparity, T maxDisp = T{0}) {
        int H = static_cast<int>(disparity.rows());
        int W = static_cast<int>(disparity.cols());
        T maxVal = std::is_floating_point_v<T> ? T{1} : T{255};

        // 最大視差を自動検出
        if (maxDisp <= T{0}) {
            for (int y = 0; y < H; ++y)
                for (int x = 0; x < W; ++x)
                    if (disparity(y, x) > maxDisp) maxDisp = disparity(y, x);
            if (maxDisp <= T{0}) maxDisp = T{1};
        }

        calx::image::Image<T> result(H, W, 3, T{0});

        for (int y = 0; y < H; ++y) {
            for (int x = 0; x < W; ++x) {
                double d = static_cast<double>(disparity(y, x));
                double norm = std::clamp(d / static_cast<double>(maxDisp), 0.0, 1.0);

                // Jet カラーマップ風
                double r, g, b;
                if (norm < 0.25) {
                    r = 0; g = 4.0 * norm; b = 1.0;
                } else if (norm < 0.5) {
                    r = 0; g = 1.0; b = 1.0 - 4.0 * (norm - 0.25);
                } else if (norm < 0.75) {
                    r = 4.0 * (norm - 0.5); g = 1.0; b = 0;
                } else {
                    r = 1.0; g = 1.0 - 4.0 * (norm - 0.75); b = 0;
                }

                result.channel(0)(y, x) = static_cast<T>(r * static_cast<double>(maxVal));
                result.channel(1)(y, x) = static_cast<T>(g * static_cast<double>(maxVal));
                result.channel(2)(y, x) = static_cast<T>(b * static_cast<double>(maxVal));
            }
        }

        return result;
    }

} // namespace calx::cv

#endif // CALX_IMAGE_CV_STEREO_MATCHING_HPP
