#include <iostream>
#include <graphene/surface_mesh/data_structure/Surface_mesh.h>
#include <opencv2/opencv.hpp>
#include <graphene/surface_mesh/algorithms/texture_processing/utility/Poisson_image_editing.h>
#include <graphene/surface_mesh/algorithms/texture_processing/Texture_neighborhood.h>
#include <graphene/surface_mesh/algorithms/texture_processing/Face_id_map.h>
#include <graphene/surface_mesh/algorithms/texture_processing/Edge_id_map.h>


bool seamless_clone(cv::Mat& image, 
                    const cv::Mat& image_src, 
                    const cv::Mat& image_mask, 
                    bool& resized, 
                    const bool use_opencv_PIE, 
                    const cv::Mat& dirichlet_mask,
                    const Texel_neighborhood& mapping_neighborhood);

void 
start_texture_merge_pie(std::string src_fn, 
                        std::string dst_fn, 
                        std::string mask_fn,
                        std::string dirichlet_mask_fn,
                        graphene::surface_mesh::Surface_mesh mesh);

int main (int argc, char** argv)
{

    std::cout << "Load all inputs for PIE" << std::endl;
    graphene::surface_mesh::Surface_mesh mesh;
    mesh.read(argv[1]);

    std::cout << "Compute PIE" << std::endl;
    if(argc == 6)
        start_texture_merge_pie(argv[2], argv[3], argv[4], argv[5], mesh);
    else if (argc == 5)
        start_texture_merge_pie(argv[2], argv[3], argv[4], "", mesh);

    std::cout << "Write result" << std::endl;
}

void 
start_texture_merge_pie(std::string src_fn, 
                             std::string dst_fn, 
                             std::string mask_fn,
                             std::string dirichlet_mask_fn,
                             graphene::surface_mesh::Surface_mesh mesh)
{

    std::cout << "Load " << src_fn << std::endl;
    cv::Mat src = cv::imread(src_fn, cv::IMREAD_COLOR);
    std::cout << "Load " << dst_fn << std::endl;
    cv::Mat dst = cv::imread(dst_fn, cv::IMREAD_COLOR);
    std::cout << "Load " << mask_fn << std::endl;
    cv::Mat mask = cv::imread(mask_fn, cv::IMREAD_GRAYSCALE);

    cv::Mat dirichlet_mask;
    if (!dirichlet_mask_fn.empty()) {
        std::cout << "Load " << dirichlet_mask_fn << std::endl;
        dirichlet_mask = cv::imread(dirichlet_mask_fn, cv::IMREAD_GRAYSCALE);
    } else {
        std::cout << "No dirichlet mask specified " << std::endl;
    }

    std::cout << "Save " << src_fn << std::endl;
    cv::imwrite("src.jpg", src);
    std::cout << "Save " << dst_fn << std::endl;
    cv::imwrite("dst.jpg", dst);
    std::cout << "Save " << mask_fn << std::endl;
    cv::imwrite("mask.jpg", mask);
    
    if (!dirichlet_mask_fn.empty()) {
        dirichlet_mask = cv::imread(dirichlet_mask_fn, cv::IMREAD_GRAYSCALE);
        std::cout << "Save " << dirichlet_mask_fn << std::endl;
        cv::imwrite("dirichlet_mask.jpg", dirichlet_mask);
    }
    

    cv::Mat face_id_tex = cv::Mat::zeros(src.size(), CV_32SC1);
    face_id_map(mesh, face_id_tex);
    cv::imwrite("face_id.jpg", face_id_tex);

    cv::Mat halfedge_id_tex = cv::Mat::zeros(src.size(), CV_32SC1);
    halfedge_id_map(mesh, halfedge_id_tex);
    cv::imwrite("halfedge_id.jpg", halfedge_id_tex);

    Texel_neighborhood resolved_neighborhood;
    resolve_uv_neighborhood(mesh, face_id_tex, halfedge_id_tex, resolved_neighborhood);


    bool resized = false;
    std::cout << "Start seamless clone" << std::endl;
    if (! seamless_clone(dst, src, mask, resized, false, dirichlet_mask, resolved_neighborhood))
    {
        std::cout << "Texture_processing::texturemerge: [ERROR] Texture copy and paste failed." << std::endl;
    }
    else {std::cout << "Seamless clone complete!" << std::endl;}
    cv::imwrite("result_merged.jpg", dst);

}



bool seamless_clone(cv::Mat& image, const cv::Mat& image_src, const cv::Mat& image_mask, bool& resized, const bool use_opencv_PIE, const cv::Mat& dirichlet_mask,
               const Texel_neighborhood& mapping_neighborhood)
{
    if ( image_mask.channels() != 1 )
    {
        std::cout << "[ERROR] image mask should have one channel only! Aborting ..." << std::endl;
        return false;
    } else {
        std::cout << "Image mask has one channel :D" << std::endl;
    }

    if( !image.data )
    {
        std::cout << "[ERROR] loading image! Aborting ..." << std::endl;
        return false;
    } else {
        std::cout << "Loaded image successfully!" << std::endl;
    }

    if( !image_src.data )
    {
        std::cout << "[ERROR] loading source image! Aborting ..." << std::endl;
        return false;
    } else {
        std::cout << "Loaded src successfully!" << std::endl;
    }

    if( !image_mask.data )
    {
        std::cout << "[ERROR] loading image mask! Aborting ..." << std::endl;
        return false;
    } else {
        std::cout << "Loaded mask successfully!" << std::endl;
    }

    // check ratios
    const unsigned int image_rows   = image.rows;
    const unsigned int image_cols   = image.cols;
    const float        image_ratio  = image_cols / image_rows;
    const cv::Size     image_size(image_cols, image_rows);

    const unsigned int image_src_rows  = image_src.rows;
    const unsigned int image_src_cols  = image_src.cols;
    const float        image_src_ratio = image_src_cols / image_src_rows;

    const unsigned int image_mask_rows  = image_mask.rows;
    const unsigned int image_mask_cols  = image_mask.cols;
    const float        image_mask_ratio = image_mask_cols / image_mask_rows;

    unsigned int dirichlet_mask_rows;
    unsigned int dirichlet_mask_cols;
    float        dirichlet_mask_ratio;
    if (dirichlet_mask.data)
    {
        dirichlet_mask_rows  = dirichlet_mask.rows;
        dirichlet_mask_cols  = dirichlet_mask.cols;
        dirichlet_mask_ratio = dirichlet_mask_cols / dirichlet_mask_rows;
    }

    if (image_src_ratio  != image_ratio ||
        image_mask_ratio != image_ratio)
    {
        std::cout << "[ERROR] images have incompatible ratios." << std::endl;
        return false;
    } else {
        std::cout << "Textures have compatible ratios!" << std::endl;
    }

    cv::Mat local_image          = image.clone();
    cv::Mat local_image_src      = image_src.clone();
    cv::Mat local_image_mask     = image_mask.clone();
    cv::Mat local_dirichlet_mask;
    if (dirichlet_mask.data)
    {
        local_dirichlet_mask = dirichlet_mask.clone();
    }

    cv::Size largest_size = image_size;
    if (local_image_src.size().width > largest_size.width)
    {
        largest_size = local_image_src.size();
    }
    if (local_image_mask.size().width > largest_size.width)
    {
        largest_size = local_image_mask.size();
    }
    if (dirichlet_mask.data)
    {
        if (local_dirichlet_mask.size().width > largest_size.width)
        {
            largest_size = local_dirichlet_mask.size();
        }
    }

    std::cout << "Resize" << std::endl;
    cv::resize(local_image, local_image, largest_size);
    cv::resize(local_image_src, local_image_src, largest_size);
    cv::resize(local_image_mask, local_image_mask, largest_size);
    if (dirichlet_mask.data)
    {
        cv::resize(local_dirichlet_mask, local_dirichlet_mask, largest_size);
    } else {
        std::cout << "No dirichlet mask!" << std::endl;
    }

    // the location of the center of the src in the dst (center of white mask region)
    std::vector<cv::Point> points;
    for (unsigned int i = 0; i < local_image_mask.rows; ++i)
    {
        for (unsigned int j = 0; j < local_image_mask.cols; ++j)
        {
            if ( local_image_mask.at<uchar>(i, j) == 255 )
            {
                points.push_back(cv::Point(j, i));
            }
        }
    }
    std::cout << "Compute bounding box" << std::endl;
    cv::Rect  bbox = cv::boundingRect(points);
    cv::Point center(bbox.x + bbox.width/2, bbox.y + bbox.height/2);


    // seamlessly clone src into dst and put the results in output
    cv::Mat seamless_clone;
    if (use_opencv_PIE)
    {
        cv::seamlessClone(local_image_src, local_image, local_image_mask, center, seamless_clone, cv::NORMAL_CLONE);
    }
    else
    {
        // cv::imwrite("test_local_image_src.jpg", local_image_src);
        // cv::imwrite("test_local_image_mask.jpg", local_image_mask);
        // cv::imwrite("test_local_image.jpg", local_image);
        // cv::imwrite("test_local_dirichlet_mask.jpg", local_dirichlet_mask);
        std::cout << "Start Poisson image editing" << std::endl;
        Poisson_image_editing poisson_image_editing(local_image_src, local_image_mask, local_image, center.x, center.y, seamless_clone, local_dirichlet_mask, mapping_neighborhood);
        // cv::imwrite("test_seamless_clone.jpg", seamless_clone);
    }

    if (largest_size != image_size)
    {
        resized = true;
    }
    else
    {
        resized = false;
    }

    image = seamless_clone.clone();

    return true;
}
