#include "sp_image.h"
#include "core/math_helper.hpp"
#include "third_party/static_block.hpp"

#include <opencv2/imgcodecs.hpp>

#include <unordered_map>

namespace {
    struct type_info {
        size_t size = {};
        int cv_type = {};
    };

    struct cv_info {
        size_t size = {};
        std::type_index type = typeid(void);

        template<typename T>
        static cv_info create() {
            return {sizeof(T), typeid(T)};
        }
    };

    template<typename Key, typename Value>
    struct map_proxy : std::unordered_map<Key, Value> {
        template<typename T>
        auto query(T &&key) const {
            const auto iter = this->find(key);
            assert(iter != this->end());
            return iter->second;
        }
    };

    using type_map_type = map_proxy<std::type_index, type_info>;
    type_map_type type_map;

    static_block {
        type_map[typeid(uchar1)] = {sizeof(uchar1), CV_8UC1};
        type_map[typeid(uchar2)] = {sizeof(uchar2), CV_8UC2};
        type_map[typeid(uchar3)] = {sizeof(uchar3), CV_8UC3};
        type_map[typeid(uchar4)] = {sizeof(uchar4), CV_8UC4};
        type_map[typeid(ushort1)] = {sizeof(ushort1), CV_16UC1};
        type_map[typeid(float1)] = {sizeof(float1), CV_32FC1};
        type_map[typeid(float2)] = {sizeof(float2), CV_32FC2};
    }

    using cv_map_type = map_proxy<int, cv_info>;
    cv_map_type cv_map;

    static_block {
        cv_map[CV_8UC1] = cv_info::create<uchar1>();
        cv_map[CV_8UC2] = cv_info::create<uchar2>();
        cv_map[CV_8UC3] = cv_info::create<uchar3>();
        cv_map[CV_8UC4] = cv_info::create<uchar4>();
        cv_map[CV_16UC1] = cv_info::create<ushort1>();
        cv_map[CV_32FC1] = cv_info::create<float1>();
        cv_map[CV_32FC2] = cv_info::create<float2>();
    }

    auto to_index_pack(const cv::Size size) {
        auto ret = index_pack<image_rank>();
        ret[0] = size.width;
        ret[1] = size.height;
        return ret;
    }
}

cv::Size sp_image::cv_size() const {
    return cv::Size(width(), height());
}

int sp_image::cv_type() const {
    return type_map.query(type).cv_type;
}

cv::Mat sp_image::cv_mat(void *ptr) const {
    return cv::Mat(cv_size(), cv_type(), start_ptr(ptr), pitch());
}

cv::cuda::GpuMat sp_image::cv_gpu_mat(void *ptr) const {
    return cv::cuda::GpuMat(cv_size(), cv_type(), start_ptr(ptr), pitch());
}

sp_image sp_image::sub_view(const cv::Size size, const cv::Size start) const {
    auto ret = *this;
    *ret.array_base() = base_type::sub_view(
        to_index_pack(size), to_index_pack(start));
    return ret;
}

sp_image sp_image::create_impl(const cv::Size size, const size_t align,
                               const std::type_index type) {
    const auto type_size = type_map.query(type).size;
    const auto pitch = alignment_round(size.width * type_size, align);
    auto ret = sp_image();
    *ret.array_base() = base_type::create(
        to_index_pack(size), pitch, type_size);
    ret.type = type;
    return ret;
}

sp_image sp_image::create_impl(const cv::Size size, const size_t align, int cv_type) {
    return create_impl(size, align, cv_map.query(cv_type).type);
}

sp_image sp_image::create_impl(const cv::Size size, const void *ptr,
                               const std::type_index type) {
    auto ret = create_impl(size, 1, type);
    const auto write_helper = write_access_helper(ret.host());
    memcpy(ret.start_ptr(write_helper.ptr()), ptr, ret.byte_size());
    return ret;
}

sp_image sp_image::create_impl(const cv::Size size, const void *ptr, int cv_type) {
    return create_impl(size, ptr, cv_map.query(cv_type).type);
}

sp_image sp_image::cast_view_impl(const std::type_index type) const {
    auto ret = *this;
    const auto type_size = type_map.query(type).size;
    *ret.array_base() = base_type::cast_view(type_size);
    ret.type = type;
    return ret;
}

sp_image sp_image::cast_view(int cv_type) const {
    return cast_view_impl(cv_map.query(cv_type).type);
}


sp_image sp_image::create(const cv::Mat &mat) {
    assert(mat.size.dims() == image_rank);
    assert(mat.isContinuous());
    return create_impl(mat.size(), mat.data,
                       cv_map.query(mat.type()).type);
}

sp_image sp_image::create_like(const sp_image &img, int cv_type) {
    if (cv_type == 0) {
        cv_type = img.cv_type();
    }
    return create(cv_type, img.cv_size());
}

sp_image sp_image::from_file(const std::string &path) {
    return create(cv::imread(path));
}

using image_ndarray_proxy = ndarray_proxy<image_rank>;
using image_index_pack = index_pack<image_rank>;

template<>
void copy_ndarray(const image_ndarray_proxy &src, image_ndarray_proxy &dst, cudaMemcpyKind kind) {
    assert(src.shape_array() == dst.shape_array());
    assert(src.byte_width() == dst.byte_width());
    if (kind == cudaMemcpyDefault) { kind = determine_copy_kind(src, dst); }

    switch (kind) {
#define TEMPLATE(src_loc, dst_loc) \
        auto access = pair_access_helper(src.src_loc(), dst.dst_loc()); \
        const auto src_ptr = src.start_ptr(access.read_ptr()); \
        const auto dst_ptr = dst.start_ptr(access.write_ptr()); \
        CUDA_API_CHECK(cudaMemcpy2DAsync( \
            dst_ptr, dst.pitch(), src_ptr, src.pitch(), \
            src.byte_width(), src.height(), kind, current_cuda_stream())); (void) 0

        //@formatter:off
        case cudaMemcpyDeviceToDevice: { TEMPLATE(cuda, cuda); break; }
        case cudaMemcpyDeviceToHost: { TEMPLATE(cuda, host); break; }
        case cudaMemcpyHostToDevice: { TEMPLATE(host, cuda); break; }
        case cudaMemcpyHostToHost: { TEMPLATE(host, host); break; }
        default: { assert(false); }
        //@formatter:on
#undef TEMPLATE
    }
}

void copy_sp_image(const sp_image &src, sp_image &dst, const cudaMemcpyKind kind) {
    assert(src.type == dst.type);
    copy_ndarray(src, dst, kind);
    dst.merge_meta(src);
}

sp_image create_dense(const sp_image &src) {
    sp_image ret = src;
    *ret.array_base() = create_dense(*src.array_base());
    return ret;
}

image_mem_info to_mem_v1(const sp_image &img, void *ptr,
                         const memory_location loc) {
    auto ret = image_mem_info();
    ret.ptr = std::shared_ptr<void>(
        img.start_ptr(ptr), [](void *) { (void) 0; });
    ret.loc = loc;
    ret.width = img.byte_width();
    ret.pitch = img.pitch();
    ret.height = img.height();
    return ret;
}