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

// optflow.hpp — Lucas-Kanade オプティカルフロー (ピラミッド LK)
//
// 新規実装:
//   - 疎オプティカルフロー (特徴点追跡)
//   - ピラミッド LK (マルチスケール)
//   - 画像勾配ベースの最小二乗

#ifndef CALX_IMAGE_CV_OPTFLOW_HPP
#define CALX_IMAGE_CV_OPTFLOW_HPP

#include <image/Image.hpp>
#include <image/pyramid.hpp>
#include <image/cv/harris.hpp>  // KeyPoint, sobelGradient
#include <cmath>
#include <vector>

namespace calx::cv {

    /// フロー結果
    struct FlowResult {
        double dx, dy;     // フロー (変位)
        bool valid;        // 追跡成功かどうか

        FlowResult() : dx(0), dy(0), valid(false) {}
        FlowResult(double dx, double dy, bool valid) : dx(dx), dy(dy), valid(valid) {}
    };

    /// 単一レベルの Lucas-Kanade (ウィンドウ内最小二乗)
    /// @param Ix, Iy, It  画像勾配
    /// @param px, py      追跡点の座標
    /// @param winSize     ウィンドウサイズ (奇数)
    /// @return            フロー (dx, dy)
    template<typename T>
    FlowResult lucasKanadeSingle(const calx::Matrix<T>& Ix,
                                 const calx::Matrix<T>& Iy,
                                 const calx::Matrix<T>& It,
                                 double px, double py,
                                 int winSize = 21) {
        int H = static_cast<int>(Ix.rows());
        int W = static_cast<int>(Ix.cols());
        int half = winSize / 2;

        int cx = static_cast<int>(std::round(px));
        int cy = static_cast<int>(std::round(py));

        if (cx < half || cx >= W - half || cy < half || cy >= H - half)
            return {};

        // 2x2 の正規方程式 A^T A d = A^T b
        double sumIxIx = 0, sumIxIy = 0, sumIyIy = 0;
        double sumIxIt = 0, sumIyIt = 0;

        for (int dy = -half; dy <= half; ++dy) {
            for (int dx = -half; dx <= half; ++dx) {
                int yy = cy + dy, xx = cx + dx;
                double ix = static_cast<double>(Ix(yy, xx));
                double iy = static_cast<double>(Iy(yy, xx));
                double it = static_cast<double>(It(yy, xx));

                sumIxIx += ix * ix;
                sumIxIy += ix * iy;
                sumIyIy += iy * iy;
                sumIxIt += ix * it;
                sumIyIt += iy * it;
            }
        }

        // det(A^T A) で逆行列が存在するか確認
        double det = sumIxIx * sumIyIy - sumIxIy * sumIxIy;
        if (std::abs(det) < 1e-10) return {};

        double invDet = 1.0 / det;
        double vx = -(sumIyIy * sumIxIt - sumIxIy * sumIyIt) * invDet;
        double vy = -(sumIxIx * sumIyIt - sumIxIy * sumIxIt) * invDet;

        return {vx, vy, true};
    }

    /// ピラミッド Lucas-Kanade オプティカルフロー
    /// @param img1       前フレーム (1ch)
    /// @param img2       現フレーム (1ch)
    /// @param points     追跡する特徴点
    /// @param maxLevel   ピラミッドの最大レベル
    /// @param winSize    LK ウィンドウサイズ
    /// @param maxIter    各レベルでの反復回数
    /// @return           各特徴点のフロー結果
    template<typename T>
    std::vector<FlowResult> pyramidLK(const image::Image<T>& img1,
                                      const image::Image<T>& img2,
                                      const std::vector<KeyPoint>& points,
                                      int maxLevel = 3,
                                      int winSize = 21,
                                      int maxIter = 30) {
        int nLevels = maxLevel + 1;
        auto pyr1 = image::buildPyramid(img1, nLevels);
        auto pyr2 = image::buildPyramid(img2, nLevels);

        // 実際のレベル数を調整
        int actualLevels = static_cast<int>(std::min(pyr1.size(), pyr2.size()));

        std::vector<FlowResult> results(points.size());

        for (std::size_t i = 0; i < points.size(); ++i) {
            double guessX = 0, guessY = 0;

            // 粗いレベルから細かいレベルへ
            for (int lv = actualLevels - 1; lv >= 0; --lv) {
                double scale = std::pow(2.0, lv);
                double px = points[i].x / scale;
                double py = points[i].y / scale;

                const auto& ch1 = pyr1[lv].channel(0);
                const auto& ch2 = pyr2[lv].channel(0);

                int H = static_cast<int>(ch1.rows());
                int W = static_cast<int>(ch1.cols());

                // 勾配計算
                auto [Ix, Iy] = sobelGradient(ch1);

                // It = I2(x + guess) - I1(x) (バイリニア補間)
                calx::Matrix<T> It(H, W, T{0});
                for (int y = 0; y < H; ++y) {
                    for (int x = 0; x < W; ++x) {
                        double sx = x + guessX;
                        double sy = y + guessY;
                        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) continue;
                        double fx = sx - x0;
                        double fy = sy - y0;

                        T val2 = static_cast<T>(
                            (1 - fy) * (1 - fx) * ch2(y0, x0) +
                            (1 - fy) * fx       * ch2(y0, x0 + 1) +
                            fy       * (1 - fx) * ch2(y0 + 1, x0) +
                            fy       * fx       * ch2(y0 + 1, x0 + 1));
                        It(y, x) = val2 - ch1(y, x);
                    }
                }

                // LK 反復
                for (int iter = 0; iter < maxIter; ++iter) {
                    auto flow = lucasKanadeSingle(Ix, Iy, It, px, py, winSize);
                    if (!flow.valid) break;

                    guessX += flow.dx;
                    guessY += flow.dy;

                    if (flow.dx * flow.dx + flow.dy * flow.dy < 0.01) break;
                }

                // 次のレベルへスケールアップ
                if (lv > 0) {
                    guessX *= 2.0;
                    guessY *= 2.0;
                }
            }

            results[i] = {guessX, guessY, true};
        }

        return results;
    }

    /// フロー結果を画像に描画
    template<typename T>
    image::Image<T> drawOpticalFlow(const image::Image<T>& img,
                                     const std::vector<KeyPoint>& points,
                                     const std::vector<FlowResult>& flows,
                                     T scale = T{1}) {
        image::Image<T> canvas = (img.numChannels() >= 3) ? img : img.toRgb();
        T maxVal = std::is_floating_point_v<T> ? T{1} : T{255};

        for (std::size_t i = 0; i < points.size() && i < flows.size(); ++i) {
            if (!flows[i].valid) continue;

            int x0 = static_cast<int>(points[i].x);
            int y0 = static_cast<int>(points[i].y);
            int x1 = x0 + static_cast<int>(flows[i].dx * static_cast<double>(scale));
            int y1 = y0 + static_cast<int>(flows[i].dy * static_cast<double>(scale));

            canvas.drawLine(x0, y0, x1, y1, {T{0}, maxVal, T{0}}, 1);
            canvas.drawCircle(x0, y0, 2, {maxVal, T{0}, T{0}}, 1);
        }

        return canvas;
    }

} // namespace calx::cv

#endif // CALX_IMAGE_CV_OPTFLOW_HPP
