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

// optflow_dense.hpp — Farneback 密オプティカルフロー
//
// 新規実装:
//   - 2次多項式展開に基づく密フロー推定
//   - ピラミッド構造によるマルチスケール対応
//   - 各レベルで反復的にフローを精密化

#ifndef CALX_IMAGE_CV_OPTFLOW_DENSE_HPP
#define CALX_IMAGE_CV_OPTFLOW_DENSE_HPP

#include <image/Image.hpp>
#include <image/pyramid.hpp>
#include <image/filters/ndimage.hpp>
#include <image/cv/harris.hpp>  // sobelGradient
#include <math/core/matrix.hpp>
#include <algorithm>
#include <cmath>
#include <vector>

namespace calx::cv {

    /// 密フロー結果 (全ピクセルの変位)
    template<typename T>
    struct DenseFlow {
        calx::Matrix<T> flowX;  // 水平方向フロー (H x W)
        calx::Matrix<T> flowY;  // 垂直方向フロー (H x W)
    };

    namespace detail_optflow {

        /// ガウシアン重みカーネルを生成
        template<typename T>
        calx::Matrix<T> gaussianWeights(int winSize) {
            calx::Matrix<T> w(winSize, winSize);
            int half = winSize / 2;
            T sigma = static_cast<T>(winSize) / T{4};
            T sum = T{0};
            for (int y = -half; y <= half; ++y) {
                for (int x = -half; x <= half; ++x) {
                    T val = std::exp(-(static_cast<T>(x * x + y * y)) / (T{2} * sigma * sigma));
                    w(y + half, x + half) = val;
                    sum += val;
                }
            }
            // 正規化
            for (int y = 0; y < winSize; ++y)
                for (int x = 0; x < winSize; ++x)
                    w(y, x) /= sum;
            return w;
        }

        /// バイリニア補間
        template<typename T>
        T bilinearSample(const calx::Matrix<T>& img, double sx, double sy) {
            int H = static_cast<int>(img.rows());
            int W = static_cast<int>(img.cols());
            int x0 = static_cast<int>(std::floor(sx));
            int y0 = static_cast<int>(std::floor(sy));
            if (x0 < 0 || x0 >= W - 1 || y0 < 0 || y0 >= H - 1)
                return T{0};
            double fx = sx - x0;
            double fy = sy - y0;
            return static_cast<T>(
                (1 - fy) * (1 - fx) * static_cast<double>(img(y0, x0)) +
                (1 - fy) * fx       * static_cast<double>(img(y0, x0 + 1)) +
                fy       * (1 - fx) * static_cast<double>(img(y0 + 1, x0)) +
                fy       * fx       * static_cast<double>(img(y0 + 1, x0 + 1)));
        }

        /// 2次多項式展開: 各ピクセルで I(x+dx,y+dy) ~ a20*dx^2 + a11*dx*dy + a02*dy^2 + a10*dx + a01*dy + a00
        /// を局所ウィンドウで最小二乗フィット。
        /// 返り値: (A, b, c) 各ピクセルにおける 2x2 行列 A、2x1 ベクトル b、スカラー c
        /// A = [[a20, a11/2], [a11/2, a02]],  b = [a10, a01]^T
        template<typename T>
        void polynomialExpansion(const calx::Matrix<T>& img,
                                 calx::Matrix<T>& A00, calx::Matrix<T>& A01, calx::Matrix<T>& A11,
                                 calx::Matrix<T>& b0,  calx::Matrix<T>& b1,
                                 int polyN, const calx::Matrix<T>& weights) {
            int H = static_cast<int>(img.rows());
            int W = static_cast<int>(img.cols());
            int half = polyN / 2;

            A00.resize(H, W); A01.resize(H, W); A11.resize(H, W);
            b0.resize(H, W);  b1.resize(H, W);

            // 各ピクセルで局所多項式フィット
            for (int y = 0; y < H; ++y) {
                for (int x = 0; x < W; ++x) {
                    // 加重最小二乗: 基底 [dx^2, dx*dy, dy^2, dx, dy, 1]
                    // 正規方程式を陽に構築 (6x6 は重いので、対称性を利用して簡略化)
                    // ここでは勾配ベースの簡略版: Sobel 近似 + Hessian 近似
                    double sGxx = 0, sGxy = 0, sGyy = 0;
                    double sGx = 0, sGy = 0;

                    for (int dy = -half; dy <= half; ++dy) {
                        for (int dx = -half; dx <= half; ++dx) {
                            int yy = y + dy, xx = x + dx;
                            if (yy < 1 || yy >= H - 1 || xx < 1 || xx >= W - 1) continue;

                            double w = static_cast<double>(weights(dy + half, dx + half));

                            // 中心差分で勾配
                            double gx = (static_cast<double>(img(yy, xx + 1)) -
                                         static_cast<double>(img(yy, xx - 1))) * 0.5;
                            double gy = (static_cast<double>(img(yy + 1, xx)) -
                                         static_cast<double>(img(yy - 1, xx))) * 0.5;

                            // 2次: Hessian の近似
                            double gxx = static_cast<double>(img(yy, xx + 1)) -
                                         2.0 * static_cast<double>(img(yy, xx)) +
                                         static_cast<double>(img(yy, xx - 1));
                            double gyy = static_cast<double>(img(yy + 1, xx)) -
                                         2.0 * static_cast<double>(img(yy, xx)) +
                                         static_cast<double>(img(yy - 1, xx));
                            double gxy = (static_cast<double>(img(yy + 1, xx + 1)) -
                                          static_cast<double>(img(yy + 1, xx - 1)) -
                                          static_cast<double>(img(yy - 1, xx + 1)) +
                                          static_cast<double>(img(yy - 1, xx - 1))) * 0.25;

                            sGxx += w * gxx;
                            sGxy += w * gxy;
                            sGyy += w * gyy;
                            sGx  += w * gx;
                            sGy  += w * gy;
                        }
                    }

                    A00(y, x) = static_cast<T>(sGxx);
                    A01(y, x) = static_cast<T>(sGxy);
                    A11(y, x) = static_cast<T>(sGyy);
                    b0(y, x)  = static_cast<T>(sGx);
                    b1(y, x)  = static_cast<T>(sGy);
                }
            }
        }

        /// フローを 2 倍にアップサンプル
        template<typename T>
        calx::Matrix<T> upsampleFlow(const calx::Matrix<T>& flow, std::size_t newH, std::size_t newW) {
            calx::Matrix<T> result(newH, newW, T{0});
            for (std::size_t y = 0; y < newH; ++y) {
                for (std::size_t x = 0; x < newW; ++x) {
                    double sy = static_cast<double>(y) * 0.5;
                    double sx = static_cast<double>(x) * 0.5;
                    result(y, x) = bilinearSample(flow, sx, sy) * T{2};
                }
            }
            return result;
        }

    } // namespace detail_optflow

    /// Farneback 密オプティカルフロー
    ///
    /// 2 次多項式展開を用いて全ピクセルの変位を推定する。
    /// ピラミッド構造で粗→細の推定を行い、各レベルで反復的に精密化する。
    ///
    /// @param prev       前フレーム (1ch グレースケール)
    /// @param next       次フレーム (1ch グレースケール)
    /// @param pyrScale   ピラミッドスケール (0.5 = 各レベルで半分)
    /// @param levels     ピラミッドレベル数
    /// @param winSize    平均化ウィンドウサイズ (奇数)
    /// @param iterations 各レベルでの反復回数
    /// @param polyN      多項式展開の近傍サイズ (5 or 7)
    /// @return           密フロー (flowX, flowY)
    template<typename T>
    DenseFlow<T> farnebackFlow(const calx::image::Image<T>& prev,
                                const calx::image::Image<T>& next,
                                T pyrScale = T{0.5}, int levels = 3,
                                int winSize = 15, int iterations = 3,
                                int polyN = 5) {
        // polyN は奇数に正規化
        if (polyN % 2 == 0) ++polyN;
        if (winSize % 2 == 0) ++winSize;

        std::size_t H = prev.height();
        std::size_t W = prev.width();

        // --- ピラミッド構築 ---
        // pyrScale に基づき各レベルのサイズを決定
        struct LevelInfo {
            calx::Matrix<T> imgPrev;
            calx::Matrix<T> imgNext;
            std::size_t h, w;
        };
        std::vector<LevelInfo> pyramid(levels);

        pyramid[0].imgPrev = prev.channel(0);
        pyramid[0].imgNext = next.channel(0);
        pyramid[0].h = H;
        pyramid[0].w = W;

        for (int lv = 1; lv < levels; ++lv) {
            std::size_t nh = static_cast<std::size_t>(pyramid[lv - 1].h * static_cast<double>(pyrScale));
            std::size_t nw = static_cast<std::size_t>(pyramid[lv - 1].w * static_cast<double>(pyrScale));
            if (nh < 4 || nw < 4) {
                levels = lv;
                pyramid.resize(levels);
                break;
            }

            // ガウシアンぼかし + ダウンサンプリング
            int kSize = 5;
            auto kernel = calx::image::gaussianKernel<T>(kSize, T{1.1});

            auto blurP = calx::image::convolve(pyramid[lv - 1].imgPrev, kernel);
            auto blurN = calx::image::convolve(pyramid[lv - 1].imgNext, kernel);

            calx::Matrix<T> downP(nh, nw);
            calx::Matrix<T> downN(nh, nw);
            for (std::size_t y = 0; y < nh; ++y) {
                for (std::size_t x = 0; x < nw; ++x) {
                    double sy = static_cast<double>(y) * static_cast<double>(pyramid[lv - 1].h) / static_cast<double>(nh);
                    double sx = static_cast<double>(x) * static_cast<double>(pyramid[lv - 1].w) / static_cast<double>(nw);
                    downP(y, x) = detail_optflow::bilinearSample(blurP, sx, sy);
                    downN(y, x) = detail_optflow::bilinearSample(blurN, sx, sy);
                }
            }

            pyramid[lv].imgPrev = std::move(downP);
            pyramid[lv].imgNext = std::move(downN);
            pyramid[lv].h = nh;
            pyramid[lv].w = nw;
        }

        // --- 粗いレベルから細かいレベルへフロー推定 ---
        calx::Matrix<T> flowX, flowY;

        auto weights = detail_optflow::gaussianWeights<T>(polyN);

        for (int lv = levels - 1; lv >= 0; --lv) {
            std::size_t lh = pyramid[lv].h;
            std::size_t lw = pyramid[lv].w;

            // フローの初期化またはアップサンプル
            if (lv == levels - 1) {
                flowX = calx::Matrix<T>(lh, lw, T{0});
                flowY = calx::Matrix<T>(lh, lw, T{0});
            } else {
                flowX = detail_optflow::upsampleFlow(flowX, lh, lw);
                flowY = detail_optflow::upsampleFlow(flowY, lh, lw);
            }

            const auto& imgP = pyramid[lv].imgPrev;
            const auto& imgN = pyramid[lv].imgNext;
            int iH = static_cast<int>(lh);
            int iW = static_cast<int>(lw);

            // 前フレームの多項式展開
            calx::Matrix<T> pA00, pA01, pA11, pB0, pB1;
            detail_optflow::polynomialExpansion(imgP, pA00, pA01, pA11, pB0, pB1, polyN, weights);

            // 反復精密化
            for (int iter = 0; iter < iterations; ++iter) {
                // 次フレームをワープして多項式展開
                calx::Matrix<T> warped(lh, lw, T{0});
                for (int y = 0; y < iH; ++y) {
                    for (int x = 0; x < iW; ++x) {
                        double sx = x + static_cast<double>(flowX(y, x));
                        double sy = y + static_cast<double>(flowY(y, x));
                        warped(y, x) = detail_optflow::bilinearSample(imgN, sx, sy);
                    }
                }

                calx::Matrix<T> nA00, nA01, nA11, nB0, nB1;
                detail_optflow::polynomialExpansion(warped, nA00, nA01, nA11, nB0, nB1, polyN, weights);

                // 各ピクセルで 2x2 系を解いてフロー更新量を求める
                // A_avg = (A_prev + A_next) / 2,  delta_b = b_next - b_prev
                // A_avg * d = -delta_b / 2  →  d = -A_avg^{-1} * delta_b / 2
                int halfW = winSize / 2;
                for (int y = 0; y < iH; ++y) {
                    for (int x = 0; x < iW; ++x) {
                        // ウィンドウ内で集約
                        double sA00 = 0, sA01 = 0, sA11 = 0;
                        double sDb0 = 0, sDb1 = 0;

                        for (int wy = -halfW; wy <= halfW; ++wy) {
                            for (int wx = -halfW; wx <= halfW; ++wx) {
                                int yy = y + wy, xx = x + wx;
                                if (yy < 0 || yy >= iH || xx < 0 || xx >= iW) continue;

                                sA00 += static_cast<double>(pA00(yy, xx) + nA00(yy, xx));
                                sA01 += static_cast<double>(pA01(yy, xx) + nA01(yy, xx));
                                sA11 += static_cast<double>(pA11(yy, xx) + nA11(yy, xx));
                                sDb0 += static_cast<double>(nB0(yy, xx) - pB0(yy, xx));
                                sDb1 += static_cast<double>(nB1(yy, xx) - pB1(yy, xx));
                            }
                        }

                        sA00 *= 0.5; sA01 *= 0.5; sA11 *= 0.5;

                        // 2x2 逆行列: [[a, b],[b, c]]^{-1} = 1/det * [[c,-b],[-b,a]]
                        double det = sA00 * sA11 - sA01 * sA01;
                        if (std::abs(det) < 1e-12) continue;

                        double invDet = 1.0 / det;
                        double dx = -( sA11 * sDb0 - sA01 * sDb1) * invDet * 0.5;
                        double dy = -(-sA01 * sDb0 + sA00 * sDb1) * invDet * 0.5;

                        // 変位が大きすぎる場合はクランプ
                        double maxDisp = static_cast<double>(winSize);
                        dx = std::clamp(dx, -maxDisp, maxDisp);
                        dy = std::clamp(dy, -maxDisp, maxDisp);

                        flowX(y, x) += static_cast<T>(dx);
                        flowY(y, x) += static_cast<T>(dy);
                    }
                }
            }
        }

        return {std::move(flowX), std::move(flowY)};
    }

    /// 密フローを HSV カラーで可視化
    /// フロー方向 → 色相、フロー大きさ → 彩度
    template<typename T>
    calx::image::Image<T> drawDenseFlow(const DenseFlow<T>& flow) {
        int H = static_cast<int>(flow.flowX.rows());
        int W = static_cast<int>(flow.flowX.cols());
        T maxVal = std::is_floating_point_v<T> ? T{1} : T{255};

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

        // 最大フロー大きさを算出
        double maxMag = 1e-6;
        for (int y = 0; y < H; ++y) {
            for (int x = 0; x < W; ++x) {
                double fx = static_cast<double>(flow.flowX(y, x));
                double fy = static_cast<double>(flow.flowY(y, x));
                double mag = std::sqrt(fx * fx + fy * fy);
                if (mag > maxMag) maxMag = mag;
            }
        }

        for (int y = 0; y < H; ++y) {
            for (int x = 0; x < W; ++x) {
                double fx = static_cast<double>(flow.flowX(y, x));
                double fy = static_cast<double>(flow.flowY(y, x));
                double mag = std::sqrt(fx * fx + fy * fy);
                double angle = std::atan2(fy, fx); // [-pi, pi]
                double hue = (angle + 3.14159265358979323846) / (2 * 3.14159265358979323846); // [0, 1]
                double sat = std::min(mag / maxMag, 1.0);

                // HSV → RGB (V=1)
                double h6 = hue * 6.0;
                int hi = static_cast<int>(h6) % 6;
                double f = h6 - static_cast<int>(h6);
                double p = 1.0 - sat;
                double q = 1.0 - sat * f;
                double t = 1.0 - sat * (1.0 - f);
                double r, g, b;
                switch (hi) {
                    case 0: r = 1; g = t; b = p; break;
                    case 1: r = q; g = 1; b = p; break;
                    case 2: r = p; g = 1; b = t; break;
                    case 3: r = p; g = q; b = 1; break;
                    case 4: r = t; g = p; b = 1; break;
                    default:r = 1; g = p; b = q; break;
                }

                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_OPTFLOW_DENSE_HPP
