#include "image_process.h"

#include <glm/ext/matrix_transform.hpp>
#include <glm/gtx/matrix_transform_2d.hpp>

#include <opencv2/cudaarithm.hpp>
#include <opencv2/cudaimgproc.hpp>
#include <opencv2/cudawarping.hpp>
#include <opencv2/imgcodecs.hpp>

namespace {
    // TODO: hack OpenCV code to make it support construction from cudaStream_t
    thread_local std::optional<cv::cuda::Stream> cv_stream;

    auto &get_cv_stream() {
        current_cuda_stream(); // initialize CUDA
        if (!cv_stream) [[unlikely]] {
            cv_stream.emplace();
        }
        return *cv_stream;
    }
}

size_t normal_height_to_nv12(const size_t height) {
    assert(height % 2 == 0);
    return height / 2 * 3;
}

size_t nv12_height_to_normal(const size_t height) {
    assert(height % 3 == 0);
    return height / 3 * 2;
}

cv::Size normal_size_to_nv12(const cv::Size size) {
    return cv::Size(size.width, normal_height_to_nv12(size.height));
}

cv::Size nv12_size_to_normal(cv::Size size) {
    return cv::Size(size.width, nv12_height_to_normal(size.height));
}

sp_image nv12_luma_view(const sp_image &img) {
    assert(img.cv_type() == CV_8UC1);
    const auto luma_size = nv12_size_to_normal(img.cv_size());
    return img.sub_view(luma_size);
}

sp_image nv12_chrome_view(const sp_image &img) {
    assert(img.cv_type() == CV_8UC1);
    const auto chroma_size = cv::Size(img.width(), img.height() / 3);
    const auto img_chrome = img.sub_view(
        chroma_size, cv::Size(0, nv12_height_to_normal(img.height())));
    return img_chrome.cast_view(CV_8UC2);
}

cv_stream_guard::cv_stream_guard() {
    push_cuda_stream((cudaStream_t) get_cv_stream().cudaPtr());
}

cv_stream_guard::~cv_stream_guard() {
    pop_cuda_stream();
}

cv::cuda::Stream &cv_stream_guard::cv_stream() {
    return get_cv_stream();
}

namespace {
    struct image_opencv_cuda_helper {
        const sp_image *read;
        sp_image *write;
        using proxy_type = auto_memory_info::cuda_proxy;
        cv_stream_guard stream_guard;
        pair_access_helper<proxy_type, proxy_type> access_helper;

        image_opencv_cuda_helper(const sp_image &src, sp_image &dst)
            : read(&src), write(&dst),
              access_helper(read->cuda(), write->cuda()) { (void) 0; }

        [[nodiscard]] cv::cuda::GpuMat input() const {
            return read->cv_gpu_mat(access_helper.read_ptr());
        }

        [[nodiscard]] cv::cuda::GpuMat output() const {
            return write->cv_gpu_mat(access_helper.write_ptr());
        }
    };
}

sp_image image_debayer(const sp_image &img) {
    assert(img.cv_type() == CV_8UC1);
    auto ret = sp_image::create<uchar3>(img.cv_size());
    const auto helper = image_opencv_cuda_helper(img, ret);
    cv::cuda::cvtColor(helper.input(), helper.output(),
                       cv::COLOR_BayerRG2BGR, 3, get_cv_stream());
    ret.merge_meta(img);
    return ret;
}

void image_resize(const sp_image &src, sp_image &dst) {
    assert(src.cv_type() == dst.cv_type());
    const auto helper = image_opencv_cuda_helper(src, dst);
    cv::cuda::resize(helper.input(), helper.output(),
                     dst.cv_size(), 0, 0, cv::INTER_LINEAR, get_cv_stream());
    dst.merge_meta(src);
}

sp_image image_resize(const sp_image &img, const cv::Size size) {
    auto ret = sp_image::create(img.cv_type(), size);
    image_resize(img, ret);
    return ret;
}

sp_image image_flip_y(const sp_image &img) {
    auto ret = sp_image::create(img.cv_type(), img.cv_size());
    const auto helper = image_opencv_cuda_helper(img, ret);
    cv::cuda::flip(helper.input(), helper.output(), 1, get_cv_stream()); // flip vertically
    ret.merge_meta(img);
    return ret;
}

sp_image image_warp_affine(const sp_image &img, const glm::mat3 &matrix) {
    auto cv_matrix = cv::Mat(2, 3, CV_32FC1);
    for (auto i = 0; i < 3; ++i)
        for (auto j = 0; j < 2; ++j) {
            cv_matrix.at<float>(j, i) = matrix[i][j];
        }
    auto ret = sp_image::create_like(img);
    const auto helper = image_opencv_cuda_helper(img, ret);
    cv::cuda::warpAffine(helper.input(), helper.output(),
                         cv_matrix, img.cv_size(), cv::INTER_LINEAR,
                         cv::BORDER_CONSTANT, {},
                         get_cv_stream());
    ret.merge_meta(img);
    return ret;
}

namespace {
    float pixel_center(const float size) {
        return 0.5f * size - 0.5f;
    }
}

sp_image image_rotate(const sp_image &img, const float angle,
                      std::optional<glm::vec2> center) {
    if (!center) {
        center = glm::vec2(pixel_center(img.width()),
                           pixel_center(img.height()));
    }
    auto matrix = glm::identity<glm::mat3>();
    matrix = glm::translate(matrix, -*center);
    matrix = glm::rotate(matrix, angle);
    matrix = glm::translate(matrix, *center);
    return image_warp_affine(img, matrix);
}

sp_image image_translate(const sp_image &img, const glm::vec2 offset) {
    const auto matrix = glm::translate(glm::identity<glm::mat3>(), offset);
    return image_warp_affine(img, matrix);
}

image_stereo_pair image_stereo_split(const sp_image &img) {
    assert(img.width() % 2 == 0);
    const auto mono_size = cv::Size(img.width() / 2, img.height());
    auto ret_left = sp_image::create(img.cv_type(), mono_size);
    const auto img_left = img.sub_view(mono_size);
    copy_sp_image(img_left, ret_left);
    auto ret_right = sp_image::create(img.cv_type(), mono_size);
    const auto img_right = img.sub_view(mono_size,
                                        cv::Size(img.width() / 2, 0));
    copy_sp_image(img_right, ret_right);
    return std::make_tuple(ret_left, ret_right);
}

image_stereo_pair image_stereo_split_view(const sp_image &img) {
    assert(img.width() % 2 == 0);
    const auto mono_size = cv::Size(img.width() / 2, img.height());
    const auto img_left = img.sub_view(mono_size);
    const auto img_right = img.sub_view(mono_size,
                                        cv::Size(img.width() / 2, 0));
    return std::make_tuple(img_left, img_right);
}

sp_image image_stereo_combine(const sp_image &left, const sp_image &right) {
    assert(left.cv_type() == right.cv_type());
    assert(left.shape_array() == right.shape_array());
    const auto stereo_shape = cv::Size(left.width() * 2, left.height());
    auto ret_img = sp_image::create(left.cv_type(), stereo_shape);
    auto [left_view, right_view] = image_stereo_split_view(ret_img);
    copy_sp_image(left, left_view);
    copy_sp_image(right, right_view);
    return ret_img;
}

#include "image_process/cuda_impl/pixel_convert.cuh"

namespace {
    template<typename Input, typename Output>
    struct image_cuda_v2_helper {
        const sp_image *read;
        sp_image *write;
        using proxy_type = auto_memory_info::cuda_proxy;
        pair_access_helper<proxy_type, proxy_type> access_helper;

        image_cuda_v2_helper(const sp_image &src, sp_image &dst)
            : read(&src), write(&dst),
              access_helper(read->cuda(), write->cuda()) { (void) 0; }

        template<typename T = Input>
        image_type_v2<T> input() {
            return to_cuda_v2(read->as_ndarray<T>(access_helper.read_ptr()));
        }

        template<typename T = Output>
        image_type_v2<T> output() {
            return to_cuda_v2(write->as_ndarray<T>(access_helper.write_ptr()));
        }
    };
}

sp_image image_rgb_to_bgr(const sp_image &img) {
    assert(img.cv_type() == CV_8UC3);
    auto ret = sp_image::create(CV_8UC3, img.cv_size());
    auto helper = image_cuda_v2_helper<uchar3, uchar3>(img, ret);
    call_cvt_rgb_bgr_u8(helper.input(), helper.output(), current_cuda_stream());
    ret.merge_meta(img);
    return ret;
}

sp_image image_rgb_to_bgra(const sp_image &img) {
    assert(img.cv_type() == CV_8UC3);
    auto ret = sp_image::create(CV_8UC4, img.cv_size());
    auto helper = image_cuda_v2_helper<uchar3, uchar4>(img, ret);
    call_cvt_rgb_bgra_u8(helper.input(), helper.output(), current_cuda_stream());
    ret.merge_meta(img);
    return ret;
}

sp_image image_rgb_to_gray(const sp_image &img) {
    assert(img.cv_type() == CV_8UC3);
    auto ret = sp_image::create<uchar1>(img.cv_size());
    const auto helper = image_opencv_cuda_helper(img, ret);
    cv::cuda::cvtColor(helper.input(), helper.output(),
                       cv::COLOR_RGB2GRAY, 1, get_cv_stream());
    ret.merge_meta(img);
    return ret;
}

sp_image image_rgb_to_nv12(const sp_image &img) {
    assert(img.cv_type() == CV_8UC3);
    auto ret = sp_image::create(CV_8UC1, normal_size_to_nv12(img.cv_size()));
    auto helper = image_cuda_v2_helper<uchar3, uchar1>(img, ret);
    call_rgb_to_nv12(helper.input(), helper.output(), current_cuda_stream());
    ret.merge_meta(img);
    return ret;
}

sp_image image_nv12_to_rgb(const sp_image &img) {
    assert(img.cv_type() == CV_8UC1);
    auto ret = sp_image::create(CV_8UC3, nv12_size_to_normal(img.cv_size()));
    auto helper = image_cuda_v2_helper<uchar1, uchar3>(img, ret);
    call_nv12_to_rgb(helper.input(), helper.output(), current_cuda_stream());
    ret.merge_meta(img);
    return ret;
}

sp_image image_yuyv_to_rgb(const sp_image &img) {
    assert(img.cv_type() == CV_8UC2);
    auto ret = sp_image::create(CV_8UC3, img.cv_size());
    auto helper = image_cuda_v2_helper<uchar2, uchar3>(img, ret);
    call_yuyv_to_rgb(helper.input(), helper.output(), current_cuda_stream());
    ret.merge_meta(img);
    return ret;
}

sp_image image_remap_np_to_tex(const sp_image &img, const float fov, const float aspect) {
    assert(img.cv_type() == CV_32FC2);
    auto ret = sp_image::create(CV_32FC2, img.cv_size());
    auto helper = image_cuda_v2_helper<float2, float2>(img, ret);
    call_np_to_tex(helper.input(), helper.output(),
                   fov, aspect, current_cuda_stream());
    ret.merge_meta(img);
    return ret;
}

namespace {
    void image_save_opencv(sp_image img, const std::string &filename) {
        if (CV_MAT_CN(img.cv_type()) == 3) {
            img = image_rgb_to_bgr(img);
        }
        const auto helper = read_access_helper(img.host());
        const auto img_mat = img.cv_mat(helper.ptr());
        cv::imwrite(filename, img_mat);
    }
}

void image_save_jpg(const sp_image &img, const std::string &filename) {
    image_save_opencv(img, fmt::format("{}.jpg", filename));
}

void image_save_png(const sp_image &img, const std::string &filename) {
    image_save_opencv(img, fmt::format("{}.png", filename));
}

#include "render/render_utility.h"

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

    void image_callback_impl() {
        const auto img = OBJ_QUERY(sp_image, conf.in_name);
        auto ret_rect = simple_rect(0, 0, conf.size.width, conf.size.height);
        ret_rect = ret_rect.fit_aspect(img.cv_size().aspectRatio());
        auto ret_img = sp_image::create(img.cv_type(), conf.size);
        ret_img.initialize_meta();
        auto ret_view = ret_img.sub_view(cv::Size(ret_rect.width, ret_rect.height),
                                         cv::Size(ret_rect.x, ret_rect.y));
        image_resize(img, ret_view);
        if (conf.flip_y) ret_img = image_flip_y(ret_img);
        OBJ_SAVE(conf.out_name, ret_img);
    }

    void image_callback(const obj_name_type _name) {
        assert(conf.in_name == _name);
        try {
            image_callback_impl();
        } catch (...) { (void) 0; }
    }

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

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

image_output_helper::image_output_helper(create_config conf)
    : pimpl(std::make_unique<impl>(conf)) {
}

image_output_helper::~image_output_helper() = default;

struct stereo_output_helper::impl {
    create_config conf;
    obj_conn_type left_conn, right_conn;

    bool left_updated = false;
    bool right_updated = false;

    void image_callback_impl() {
        const auto left_img = OBJ_QUERY(sp_image, conf.left_name);
        const auto right_img = OBJ_QUERY(sp_image, conf.right_name);
        assert(left_img.cv_type() == right_img.cv_type());
        assert(left_img.cv_size() == right_img.cv_size());
        auto ret_size = conf.size;
        if (ret_size.empty()) {
            if (conf.halve_width) {
                ret_size = left_img.cv_size();
            } else {
                ret_size = cv::Size(left_img.width() * 2, left_img.height());
            }
        }
        assert(ret_size.width % 2 == 0);
        auto ret_rect = simple_rect(0, 0,
                                    conf.halve_width ? ret_size.width : (ret_size.width / 2),
                                    ret_size.height);
        ret_rect = ret_rect.fit_aspect(left_img.cv_size().aspectRatio());
        if (conf.halve_width) {
            ret_rect.x /= 2;
            ret_rect.width /= 2;
        }
        auto ret_img = sp_image::create(left_img.cv_type(), ret_size);
        ret_img.initialize_meta();
        auto left_view = ret_img.sub_view(cv::Size(ret_rect.width, ret_rect.height),
                                          cv::Size(ret_rect.x, ret_rect.y));
        image_resize(left_img, left_view);
        auto right_view = ret_img.sub_view(cv::Size(ret_rect.width, ret_rect.height),
                                           cv::Size(ret_rect.x + ret_size.width / 2, ret_rect.y));
        image_resize(right_img, right_view);
        if (conf.flip_y) ret_img = image_flip_y(ret_img);
        OBJ_SAVE(conf.out_name, ret_img);
    }

    void image_callback(const obj_name_type name) {
        if (name == conf.left_name) left_updated = true;
        if (name == conf.right_name) right_updated = true;
        if (!left_updated || !right_updated) return;
        try {
            image_callback_impl();
        } catch (...) { (void) 0; }
        left_updated = false;
        right_updated = false;
    }

    explicit impl(const create_config &_conf) : conf(_conf) {
        left_conn = OBJ_SIG(conf.left_name)->connect(
            [this](auto name) { image_callback(name); });
        right_conn = OBJ_SIG(conf.right_name)->connect(
            [this](auto name) { image_callback(name); });
    }

    ~impl() {
        left_conn.disconnect();
        right_conn.disconnect();
    }
};

stereo_output_helper::stereo_output_helper(create_config conf)
    : pimpl(std::make_unique<impl>(conf)) {
}

stereo_output_helper::~stereo_output_helper() = default;