#include "video_stabilization.h"
#include "image_process.h"
#include "third_party/scope_guard.hpp"

#include <opencv2/calib3d.hpp>
#include <opencv2/cudaoptflow.hpp>
#include <opencv2/cudaimgproc.hpp>

#include <Eigen/Dense>

#include <cmath>
#include <deque>
#include <ranges>
#include <vector>

namespace {
    template<typename T>
    std::vector<T> download(const cv::cuda::GpuMat &d_mat, int type) {
        auto ret = std::vector<T>(d_mat.cols);
        auto mat = cv::Mat(1, d_mat.cols, type, ret.data());
        d_mat.download(mat);
        return ret;
    }

    struct point_tracker {
        cv::Ptr<cv::cuda::CornersDetector> detector;
        cv::Ptr<cv::cuda::SparsePyrLKOpticalFlow> tracker;
        sp_image prev;

        using points_type = std::vector<cv::Point2f>;

        struct result_type {
            points_type prev_pts;
            points_type curr_pts;
        };

        using create_config = video_stabilization::point_tracker_config;
        create_config conf;

        explicit point_tracker(const create_config _conf)
            : conf(_conf) {
            detector = cv::cuda::createGoodFeaturesToTrackDetector(
                CV_8UC1, conf.max_feature_points,
                0.3, 7, 7);
            tracker = cv::cuda::SparsePyrLKOpticalFlow::create();
        }

        // Input RGB image
        std::optional<result_type> process(const sp_image &img) {
            auto closer = sg::make_scope_guard([&] { prev = img; });
            if (prev.empty()) return {};

            auto stream_guard = cv_stream_guard();
            cv::cuda::GpuMat d_prev_pts;
            if (true) {
                const auto prev_gray = image_rgb_to_gray(prev);
                const auto helper = read_access_helper(prev_gray.cuda());
                detector->detect(prev_gray.cv_gpu_mat(helper.ptr()), d_prev_pts,
                                 cv::noArray(), stream_guard.cv_stream());
            }

            cv::cuda::GpuMat d_curr_pts, d_status;
            if (true) {
                const auto prev_helper = read_access_helper(prev.cuda());
                const auto curr_helper = read_access_helper(img.cuda());
                tracker->calc(prev.cv_gpu_mat(prev_helper.ptr()),
                              img.cv_gpu_mat(curr_helper.ptr()),
                              d_prev_pts, d_curr_pts, d_status);
            }

            auto status = download<uchar>(d_status, CV_8UC1);
            auto valid_indices = std::views::iota(0ull, status.size())
                                 | std::views::filter([&](auto i) { return status[i]; });
            auto prev_pts = download<cv::Point2f>(d_prev_pts, CV_32FC2);
            auto curr_pts = download<cv::Point2f>(d_curr_pts, CV_32FC2);
            auto valid_prev_pts = valid_indices
                                  | std::views::transform([&](auto i) { return prev_pts[i]; });
            auto valid_curr_pts = valid_indices
                                  | std::views::transform([&](auto i) { return curr_pts[i]; });
            auto ret = result_type();
            std::ranges::copy(valid_prev_pts, std::back_inserter(ret.prev_pts));
            std::ranges::copy(valid_curr_pts, std::back_inserter(ret.curr_pts));
            return ret;
        }
    };

    struct transform_stabilization {
        using matrix_type = Eigen::Matrix<float, 2, 3>;
        using matrix_param_type = Eigen::Vector4f;
        using record_type = Eigen::VectorXf;
        using trajectory_type = Eigen::Matrix4Xf;

        static record_type optimize_simple(const record_type &traj) {
            const auto num = traj.rows();
            auto ret = record_type(num);
            ret.setConstant(traj.mean());
            return ret;
        }

        static trajectory_type optimize_trajectory(const trajectory_type &traj) {
            const auto num_dims = traj.rows();
            auto ret = trajectory_type(num_dims, traj.cols());
            for (auto i = 0; i < num_dims; ++i) {
                ret.row(i) = optimize_simple(traj.row(i));
            }
            return ret;
        }

        static matrix_param_type decompose_transform(const matrix_type &mat) {
            const auto scale = sqrtf(mat.block<2, 2>(0, 0).determinant());
            const auto angle = atan2f(mat(1, 0), mat(0, 0));
            const auto tx = mat(0, 2);
            const auto ty = mat(1, 2);
            return {logf(scale), angle, tx, ty};
        }

        static matrix_type compose_transform(const matrix_param_type &param) {
            const auto scale = expf(param(0)), angle = param(1);
            const auto tx = param(2), ty = param(3);
            auto ret = matrix_type();
            ret << cosf(angle), -sinf(angle), tx,
                    sinf(angle), cosf(angle), ty;
            ret.block<2, 2>(0, 0) *= scale;
            return ret;
        }

        using create_config = video_stabilization::transform_stabilization_config;
        create_config conf;

        size_t buffer_length = 0;
        std::deque<matrix_param_type> transform_buf;

        explicit transform_stabilization(const create_config &_conf)
            : conf(_conf) {
            buffer_length = 2 * conf.future_frames + 1;
        }

        std::optional<matrix_type> process(const matrix_type &mat) {
            const auto param = decompose_transform(mat);
            transform_buf.push_back(param);
            if (transform_buf.size() <= conf.future_frames) return {};
            if (transform_buf.size() > buffer_length) {
                transform_buf.pop_front();
            }

            auto curr_traj = trajectory_type(4, transform_buf.size());
            for (int i = 0; i < transform_buf.size(); i++) {
                curr_traj.col(i) = transform_buf[i];
                if (i > 0) {
                    curr_traj.col(i) += curr_traj.col(i - 1);
                }
            }

            const auto traj_opt = optimize_trajectory(curr_traj);
            const auto curr_pos = conf.future_frames;
            const auto diff = traj_opt.col(curr_pos) - curr_traj.col(curr_pos);
            const auto curr = transform_buf[curr_pos] + diff;
            return compose_transform(curr);
        }
    };
}

struct video_stabilization::impl {
    using matrix_type = transform_stabilization::matrix_type;

    static matrix_type combine_transform(const matrix_type &mat1, const matrix_type &mat2) {
        auto ret = matrix_type();
        const auto r1_mat = mat1.leftCols<2>();
        ret.leftCols<2>() = r1_mat * mat2.leftCols<2>();
        ret.rightCols<1>() = r1_mat * mat2.rightCols<1>() + mat1.rightCols<1>();
        return ret;
    }

    static matrix_type identity_transform() {
        auto ret = matrix_type();
        ret.setZero();
        ret.leftCols<2>().setIdentity();
        return ret;
    }

    static glm::mat3 to_mat3(const matrix_type &mat) {
        auto ret = glm::mat3();
        for (auto i = 0; i < 2; ++i)
            for (auto j = 0; j < 3; ++j) {
                ret[j][i] = mat(i, j);
            }
        return ret;
    }

    static matrix_type to_transform(const cv::Mat &mat) {
        assert(mat.rows == 2 && mat.cols == 3);
        assert(mat.type() == CV_64FC1);
        auto ret = matrix_type();
        for (auto i = 0; i < 2; ++i)
            for (auto j = 0; j < 3; ++j) {
                ret(i, j) = mat.at<double>(i, j);
            }
        return ret;
    }

    static constexpr auto min_tracked_points = 10;

    create_config conf;
    std::optional<point_tracker> track_left, track_right;
    std::optional<transform_stabilization> transform_stab;
    std::deque<sp_image> img_buf;

    [[nodiscard]] sp_image transform_image(const sp_image &img, matrix_type mat) const {
        const auto center = cv::Point2f(img.width() / 2.f, img.height() / 2.f);
        const auto fix_mat = cv::getRotationMatrix2D(center, 0, 1.f + conf.enlarge_factor);
        mat = combine_transform(to_transform(fix_mat), mat);
        return image_warp_affine(img, to_mat3(mat));
    }

    std::optional<sp_image> process_mono(const sp_image &img) {
        assert(!conf.stereo_mode);
        const auto pts = track_left->process(img);
        if (!pts) return {};
        matrix_type mat;
        if (pts->prev_pts.size() < min_tracked_points) [[unlikely]] {
            mat = identity_transform();
        } else {
            const auto cv_mat = cv::estimateAffinePartial2D(
                pts->prev_pts, pts->curr_pts);
            mat = to_transform(cv_mat);
        }

        const auto cur_mat = transform_stab->process(mat);
        img_buf.push_back(img);
        if (!cur_mat) return {};

        const auto cur_img = img_buf.front();
        img_buf.pop_front();
        return transform_image(cur_img, *cur_mat);
    }

    std::optional<sp_image> process_stereo(const sp_image &img) {
        assert(conf.stereo_mode);
        auto [img_left, img_right] = image_stereo_split_view(img);
        auto pts_left = track_left->process(img_left);
        auto pts_right = track_right->process(img_right);
        if (!pts_left || !pts_right) return {};

        auto num_points = pts_left->prev_pts.size() + pts_right->prev_pts.size();
        matrix_type mat;
        if (num_points < min_tracked_points) [[unlikely]] {
            mat = identity_transform();
        } else {
            auto prev_pts = std::move(pts_left->prev_pts);
            prev_pts.reserve(num_points);
            std::ranges::copy(pts_right->prev_pts, std::back_inserter(prev_pts));
            auto curr_pts = std::move(pts_left->curr_pts);
            curr_pts.reserve(num_points);
            std::ranges::copy(pts_right->curr_pts, std::back_inserter(curr_pts));
            const auto cv_mat = cv::estimateAffinePartial2D(prev_pts, curr_pts);
            mat = to_transform(cv_mat);
        }

        const auto cur_mat = transform_stab->process(mat);
        img_buf.push_back(img);
        if (!cur_mat) return {};

        const auto cur_img = img_buf.front();
        img_buf.pop_front();
        auto [cur_left, cur_right] = image_stereo_split_view(cur_img);
        cur_left = transform_image(cur_left, *cur_mat);
        cur_right = transform_image(cur_right, *cur_mat);
        return image_stereo_combine(cur_left, cur_right);
    }

    auto process(const sp_image &img) {
        if (conf.stereo_mode) {
            return process_stereo(img);
        } else {
            return process_mono(img);
        }
    }

    explicit impl(const create_config &_conf) : conf(_conf) {
        track_left.emplace(conf.pt_conf);
        if (conf.stereo_mode) {
            track_right.emplace(conf.pt_conf);
        }
        transform_stab.emplace(conf.st_conf);
    }
};

video_stabilization::video_stabilization(create_config conf)
    : pimpl(std::make_unique<impl>(conf)) { (void) 0; }

video_stabilization::~video_stabilization() = default;

std::optional<sp_image> video_stabilization::process(const sp_image &img) const {
    return pimpl->process(img);
}

#include "core/imgui_utility.hpp"

struct video_stabilization_ui::impl {
    create_config conf;
    obj_conn_type conn;

    bool passthrough = false;
    std::optional<BS::thread_pool> work_tp;
    std::optional<video_stabilization> video_stab;

    void image_callback_impl() const {
        if (!conf.opts.stereo_mode) {
            const auto img = OBJ_QUERY(sp_image, conf.in_name);
            if (const auto ret = video_stab->process(img)) {
                OBJ_SAVE(conf.out_name, *ret);
            }
        } else {
            auto img_left = OBJ_QUERY(sp_image, conf.in_name);
            auto img_right = OBJ_QUERY(sp_image, conf.in2_name);
            auto img = image_stereo_combine(img_left, img_right);
            if (auto ret = video_stab->process(img)) {
                auto [img_left, img_right] = image_stereo_split_view(*ret);
                OBJ_SAVE(conf.out_name, img_left);
                OBJ_SAVE(conf.out2_name, img_right);
            }
        }
    }

    void image_callback() {
        if (passthrough) {
            try {
                OBJ_SAVE(conf.out_name, OBJ_QUERY(sp_image, conf.in_name));
                if (conf.opts.stereo_mode) {
                    OBJ_SAVE(conf.out2_name, OBJ_QUERY(sp_image, conf.in2_name));
                }
            } catch (...) { (void) 0; }
            return;
        }

        if (!video_stab) [[unlikely]] {
            video_stab.emplace(conf.opts);
        }
        auto task = [this] {
            try {
                image_callback_impl();
            } catch (...) { (void) 0; }
        };
        if (work_tp->get_tasks_queued() <= 1) {
            work_tp->detach_task(task);
        } else {
            SPDLOG_WARN("Too many frames for stabilization, frame dropped.");
        }
    }

    void reset_video_stab() {
        MAIN_DETACH([this] {
            work_tp->wait();
            video_stab.reset();
        });
    }

    void show_ui() {
        if (ImGui::Checkbox("Passthrough", &passthrough)) {
            if (!passthrough) {
                reset_video_stab();
            }
        }
        if (!passthrough) {
            ImGui::SameLine();
            if (ImGui::Button("Reset")) {
                reset_video_stab();
            }
            ImGui::DragInt("Frame Delay",
                           &conf.opts.st_conf.future_frames, 1, 1);
        }
    }

    explicit impl(const create_config &_conf) : conf(_conf) {
        work_tp.emplace(1);
        conn = OBJ_SIG(conf.in_name)->connect(
            [this](auto _) { image_callback(); });
    }

    ~impl() {
        conn.disconnect();
    }
};

video_stabilization_ui::video_stabilization_ui(create_config conf)
    : pimpl(std::make_unique<impl>(conf)) { (void) 0; }

video_stabilization_ui::~video_stabilization_ui() = default;

void video_stabilization_ui::show_ui() const {
    pimpl->show_ui();
}