//== INCLUDES ===================================================================


#include "Image_smooth_laplace_zero.h"

#include <vector>
#include <climits>

#ifdef HAVE_CHOLMOD
#include <Eigen/CholmodSupport>
#define SparseSPDSolver CholmodSupernodalLLT
#else
#define SparseSPDSolver SimplicialLDLT
#endif

//== IMPLEMENTATION ============================================================


Image_smooth_laplace_zero::
Image_smooth_laplace_zero(const cv::Mat& source_image,
                          cv::Mat& source_mask,
                          cv::Mat& result_image,
                          const Texel_neighborhood& neighborhood)
{
    cv::Mat src = source_image.clone();
    src.convertTo(src, CV_64FC3);

    // split the channels for the source and target image
    std::vector<cv::Mat> src_bgr;
    cv::split(src, src_bgr);
    cv::Mat src_b = src_bgr[0];
    cv::Mat src_g = src_bgr[1];
    cv::Mat src_r = src_bgr[2];

    std::vector<cv::Mat> result_bgr;

    // segment mask to: outside, boundary, inside
    cv::Mat segmentation;
    segment_mask(source_mask, segmentation, neighborhood);

    // get inner coordinated (these pixels should be changed)
    std::vector<cv::Point> inner_coordinates;
    get_inner_coordinates(segmentation, inner_coordinates);

    // get mapping from pixel coordinate to index of inner coordinates
    std::map< unsigned int, std::map<unsigned int, unsigned int> > mapping_pixelcoordinate_to_idx;
    get_mapping_pixelcoordinate_to_idx(segmentation, mapping_pixelcoordinate_to_idx);

    // assemble A
    Image_smooth_laplace_zero::SpMat A_eigen = assemble_A(segmentation, inner_coordinates, mapping_pixelcoordinate_to_idx, neighborhood);

    // blend each channel separately

    // assemble b_b
    Eigen::VectorXd b_b_eigen = assemble_b(segmentation, inner_coordinates, src_b, neighborhood);

    // assemble b_g
    Eigen::VectorXd b_g_eigen = assemble_b(segmentation, inner_coordinates, src_g, neighborhood);

    // assemble b_r
    Eigen::VectorXd b_r_eigen = assemble_b(segmentation, inner_coordinates, src_r, neighborhood);

    // solve for new pixel colors
    Eigen::VectorXd x_b;
    Eigen::VectorXd x_g;
    Eigen::VectorXd x_r;
    const bool A_is_symmetric = !valid(neighborhood);
    if (! solve_sle(A_eigen, b_b_eigen, b_g_eigen, b_r_eigen, x_b, x_g, x_r, A_is_symmetric) )
    {
        std::cerr << "Can't solve for new pixels. Aborting..." << std::endl;
        return;
    }

    cv::Mat result_b(src_b.size(), CV_64FC1);
    cv::Mat result_g(src_g.size(), CV_64FC1);
    cv::Mat result_r(src_r.size(), CV_64FC1);

    solution_to_cvMat(segmentation, inner_coordinates, x_b, x_g, x_r, result_b, result_g, result_r, src_b, src_g, src_r);

    // merge results for each channel
    result_bgr.push_back(result_b);
    result_bgr.push_back(result_g);
    result_bgr.push_back(result_r);
    cv::Mat result;
    cv::merge(result_bgr, result);

    result.convertTo(result, CV_8UC3);
    result_image = result.clone();
}

//------------------------------------------------------------------------------


void
Image_smooth_laplace_zero::
segment_mask(const cv::Mat& source_mask, cv::Mat& segmentation, const Texel_neighborhood& neighborhood) const
{
    // initialize to outside
    segmentation = cv::Mat(source_mask.size(), CV_8UC1, cv::Scalar(outside));

    // now set pixels to inside w.r.t. source mask
    for (int i = 1; i < source_mask.rows - 1; i++)
    {
        for (int j = 1; j < source_mask.cols - 1; j++)
        {
            if ( source_mask.at<uchar>(i, j) == 255 )
            {
                segmentation.at<uchar>(i, j) = inside;
            }
        }
    }

    if (!valid(neighborhood))
    {
        // Assume standard 4 neighborhood everywhere

        // and eventually update outer pixels to boundary
        for (int i = 0; i < source_mask.rows; i++)
        {
            for (int j = 0; j < source_mask.cols; j++)
            {
                if (segmentation.at<uchar>(i, j) == outside)
                {
                    //check whether we access pixels that are outside of the image
                    if (i-1 < 0 || i+1 >= source_mask.rows || j-1 < 0 || j+1 >= source_mask.cols)
                    {
                        continue;
                    }
                    if (source_mask.at<uchar>(i-1, j-1) == 255 ||
                        source_mask.at<uchar>(i-1, j)   == 255 ||
                        source_mask.at<uchar>(i-1, j+1) == 255 ||
                        source_mask.at<uchar>(i, j-1)   == 255 ||
                        source_mask.at<uchar>(i, j+1)   == 255 ||
                        source_mask.at<uchar>(i+1, j-1) == 255 ||
                        source_mask.at<uchar>(i+1, j)   == 255 ||
                        source_mask.at<uchar>(i+1, j+1) == 255)
                    {
                        segmentation.at<uchar>(i, j) = boundary;
                    }
                }
            }
        }
    }
    else
    {
        // Neighborhood is defined by passed vector

        // and eventually update outer pixels to boundary
        for (int i = 0; i < source_mask.rows; i++)
        {
            for (int j = 0; j < source_mask.cols; j++)
            {
                if (segmentation.at<uchar>(i, j) == inside)
                {
                    //check whether we access pixels that are outside of the image
                    if (i-1 < 0 || i+1 >= source_mask.rows || j-1 < 0 || j+1 >= source_mask.cols)
                    {
                        continue;
                    }

                    auto r = right(neighborhood, i, j);
                    auto l = left(neighborhood, i, j);
                    auto b = bottom(neighborhood, i, j);
                    auto t = top(neighborhood, i, j);

                    // Any inside texels whose neighbor is outside is a boundary texel
                    if (segmentation.at<uchar>(r) == outside) segmentation.at<uchar>(r) = boundary;
                    if (segmentation.at<uchar>(l) == outside) segmentation.at<uchar>(l) = boundary;
                    if (segmentation.at<uchar>(b) == outside) segmentation.at<uchar>(b) = boundary;
                    if (segmentation.at<uchar>(t) == outside) segmentation.at<uchar>(t) = boundary;
                }
            }
        }
    }

    // dirichlet boundary conditions
    for (int i = 0; i < segmentation.rows; i++)
    {
        for (int j = 0; j < segmentation.cols; j++)
        {
            if (segmentation.at<uchar>(i, j) == boundary)
            {
                segmentation.at<uchar>(i, j) = dirichlet;
            }
        }
    }

    //cv::namedWindow("Segmentation", cv::WINDOW_NORMAL);
    //cv::imshow("Segmentation", segmentation * 80);
    //cv::waitKey(0);
}


//------------------------------------------------------------------------------


void
Image_smooth_laplace_zero::
get_inner_coordinates(const cv::Mat& segmentation, std::vector<cv::Point>& result) const
{
    result.clear();

    for (int i = 0; i < segmentation.rows; i++)
    {
        for (int j = 0; j < segmentation.cols; j++)
        {
            if ( segmentation.at<uchar>(i, j) == inside )
            {
                result.push_back(cv::Point(j, i));
            }
        }
    }
}


//------------------------------------------------------------------------------


void
Image_smooth_laplace_zero::
get_mapping_pixelcoordinate_to_idx(const cv::Mat& segmentation, std::map< unsigned int, std::map<unsigned int, unsigned int> >& mapping) const
{
    unsigned int idx = 0;
    for (int i = 0; i < segmentation.rows; i++)
    {
        for (int j = 0; j < segmentation.cols; j++)
        {
            if ( segmentation.at<uchar>(i, j) == inside )
            {
                mapping[i][j] = idx;
                idx++;
            }
        }
    }
}


//------------------------------------------------------------------------------


Image_smooth_laplace_zero::SpMat
Image_smooth_laplace_zero::
assemble_A(const cv::Mat& segmentation, const std::vector<cv::Point>& inner_coordinates, const std::map< unsigned int, std::map<unsigned int, unsigned int> >& mapping_pixelcoordinate_to_idx,
           const Texel_neighborhood& neighborhood) const
{
    const unsigned int dofs = inner_coordinates.size();

    Image_smooth_laplace_zero::SpMat A(dofs, dofs);
    std::vector<Tripl> coeffs;
    coeffs.reserve( 5*inner_coordinates.size() );

#if _OPENMP >= 201310
#pragma omp declare reduction (merge : std::vector<Tripl> : omp_out.insert(omp_out.end(), omp_in.begin(), omp_in.end()))
#pragma omp parallel for reduction(merge: coeffs)
#endif
    for (unsigned int i = 0; i < inner_coordinates.size(); ++i)
    {
        const unsigned int x = inner_coordinates[i].x;
        const unsigned int y = inner_coordinates[i].y;

        unsigned int lx = x-1;
        unsigned int ly = y;
        unsigned int tx = x;
        unsigned int ty = y-1;
        unsigned int rx = x+1;
        unsigned int ry = y;
        unsigned int bx = x;
        unsigned int by = y+1;

        if (valid(neighborhood))
        {
            if (!valid(neighborhood, y, x))
            {
                printf("[ERROR] No neighborhood resolved for (%d,%d)\n", y, x);
                continue;
            }

            auto r = right(neighborhood, y, x);
            auto l = left(neighborhood, y, x);
            auto b = bottom(neighborhood, y, x);
            auto t = top(neighborhood, y, x);
            lx = l.x;
            ly = l.y;
            tx = t.x;
            ty = t.y;
            rx = r.x;
            ry = r.y;
            bx = b.x;
            by = b.y;
        }

        if (segmentation.at<uchar>(ty, tx) == inside && segmentation.at<uchar>(y-1, x) != dirichlet) // T
        {
            const unsigned int idx = mapping_pixelcoordinate_to_idx.at(ty).at(tx);
            coeffs.push_back( Tripl(i, idx, -1.0) );
        }

        if (segmentation.at<uchar>(by, bx) == inside && segmentation.at<uchar>(y+1, x) != dirichlet) // B
        {
            const unsigned int idx = mapping_pixelcoordinate_to_idx.at(by).at(bx);
            coeffs.push_back( Tripl(i, idx, -1.0) );
        }

        if (segmentation.at<uchar>(ly, lx) == inside && segmentation.at<uchar>(y, x-1) != dirichlet) // L
        {
            const unsigned int idx = mapping_pixelcoordinate_to_idx.at(ly).at(lx);
            coeffs.push_back( Tripl(i, idx, -1.0) );
        }

        if (segmentation.at<uchar>(ry, rx) == inside && segmentation.at<uchar>(y, x+1) != dirichlet) // R
        {
            const unsigned int idx = mapping_pixelcoordinate_to_idx.at(ry).at(rx);
            coeffs.push_back( Tripl(i, idx, -1.0) );
        }

        coeffs.push_back( Tripl(i, i, 4.0) );
    }

    A.setFromTriplets(coeffs.begin(), coeffs.end());

    return A;
}


//------------------------------------------------------------------------------


Eigen::VectorXd
Image_smooth_laplace_zero::
assemble_b(const cv::Mat& segmentation, const std::vector<cv::Point>& inner_coordinates, const cv::Mat& src_channel,
           const Texel_neighborhood& neighborhood) const
{
    const unsigned int dofs = inner_coordinates.size();

    Eigen::VectorXd b(dofs);
    for (unsigned int i = 0; i < inner_coordinates.size(); ++i)
    {
        const unsigned int x = inner_coordinates[i].x;
        const unsigned int y = inner_coordinates[i].y;

        unsigned int lx = x-1;
        unsigned int ly = y;
        unsigned int tx = x;
        unsigned int ty = y-1;
        unsigned int rx = x+1;
        unsigned int ry = y;
        unsigned int bx = x;
        unsigned int by = y+1;

        if (valid(neighborhood))
        {
            if (!valid(neighborhood, y, x))
            {
                printf("[ERROR] No neighborhood resolved for (%d, %d)\n", y, x);
                continue;
            }

            auto r = right(neighborhood, y, x);
            auto l = left(neighborhood, y, x);
            auto b = bottom(neighborhood, y, x);
            auto t = top(neighborhood, y, x);
            lx = l.x;
            ly = l.y;
            tx = t.x;
            ty = t.y;
            rx = r.x;
            ry = r.y;
            bx = b.x;
            by = b.y;
        }

        b(i) = 0;

        // T
        if (segmentation.at<uchar>(ty, tx) == dirichlet)
        {
            b(i) += src_channel.at<double>(ty, tx);
        }

        // B
        if (segmentation.at<uchar>(by, bx) == dirichlet)
        {
            b(i) += src_channel.at<double>(by, bx);
        }

        // L
        if (segmentation.at<uchar>(ly, lx) == dirichlet)
        {
            b(i) += src_channel.at<double>(ly, lx);
        }

        // R
        if (segmentation.at<uchar>(ry, rx) == dirichlet)
        {
           b(i) += src_channel.at<double>(ry, rx);
        }
    }

    return b;
}


//------------------------------------------------------------------------------


void
Image_smooth_laplace_zero::
solution_to_cvMat(const cv::Mat& segmentation,
                  const std::vector<cv::Point>& inner_coordinates,
                  const Eigen::VectorXd& x_eigen_b,
                  const Eigen::VectorXd& x_eigen_g,
                  const Eigen::VectorXd& x_eigen_r,
                  cv::Mat& result_channel_b,
                  cv::Mat& result_channel_g,
                  cv::Mat& result_channel_r,
                  const cv::Mat& target_channel_b,
                  const cv::Mat& target_channel_g,
                  const cv::Mat& target_channel_r) const
{
    result_channel_b = target_channel_b.clone();
    result_channel_g = target_channel_g.clone();
    result_channel_r = target_channel_r.clone();

    for (unsigned int i = 0; i < inner_coordinates.size(); ++i)
    {
        const unsigned int x = inner_coordinates[i].x;
        const unsigned int y = inner_coordinates[i].y;

        result_channel_b.at<double>(y, x) = x_eigen_b(i);
        result_channel_g.at<double>(y, x) = x_eigen_g(i);
        result_channel_r.at<double>(y, x) = x_eigen_r(i);
    }
}


//------------------------------------------------------------------------------


const bool
Image_smooth_laplace_zero::
solve_sle(const SpMat& A,
          const Eigen::VectorXd& b_b,
          const Eigen::VectorXd& b_g,
          const Eigen::VectorXd& b_r,
          Eigen::VectorXd& x_b,
          Eigen::VectorXd& x_g,
          Eigen::VectorXd& x_r,
          const bool is_symmetric) const
{
    if (is_symmetric)
    {
        // Eigen::SimplicialLDLT< SpMat > solver;
        Eigen::SparseSPDSolver< SpMat > solver;

        solver.compute(A);
        if(solver.info() != Eigen::Success)
        {
            // decomposition failed
            std::cerr << "[ERROR] Decomposition failed" << std::endl;
            return false;
        }

        x_b = solver.solve(b_b);
        if(solver.info() != Eigen::Success)
        {
            // solving failed
            std::cerr << "[ERROR] Solving for b-channel failed" << std::endl;
            return false;
        }

        x_g = solver.solve(b_g);
        if(solver.info() != Eigen::Success)
        {
            // solving failed
            std::cerr << "[ERROR] Solving for g-channel failed" << std::endl;
            return false;
        }

        x_r = solver.solve(b_r);
        if(solver.info() != Eigen::Success)
        {
            // solving failed
            std::cerr << "[ERROR] Solving for r-channel failed" << std::endl;
            return false;
        }
    }
    else
    {
        Eigen::SparseLU< SpMat > solver;

        solver.compute(A);
        if(solver.info() != Eigen::Success)
        {
            // decomposition failed
            std::cerr << "[ERROR] Decomposition failed" << std::endl;
            return false;
        }

        x_b = solver.solve(b_b);
        if(solver.info() != Eigen::Success)
        {
            // solving failed
            std::cerr << "[ERROR] Solving for b-channel failed" << std::endl;
            return false;
        }

        x_g = solver.solve(b_g);
        if(solver.info() != Eigen::Success)
        {
            // solving failed
            std::cerr << "[ERROR] Solving for g-channel failed" << std::endl;
            return false;
        }

        x_r = solver.solve(b_r);
        if(solver.info() != Eigen::Success)
        {
            // solving failed
            std::cerr << "[ERROR] Solving for r-channel failed" << std::endl;
            return false;
        }
    }

    return true;
}


//==============================================================================
