#pragma once

#include <vector>
#include <set>

#include <opencv2/core.hpp>

#include <graphene/surface_mesh/data_structure/Surface_mesh.h>
#include <graphene/geometry/distance_point_triangle.h>
#include <graphene/geometry/bary_coord.h>

#include "Face_id_map.h"

using graphene::surface_mesh::Surface_mesh;
using graphene::Point;


//
// API
//


struct Texel_neighborhood;

//
// Resolves neighborhood on the uv layout of the passed mesh.
//
// For face_id_tex and halfedge_id_tex parameters, see Face_id_map.h and Edge_id_map.h
//
// Builds up texel neighborhood that contains either 0 (boundary texels) or 4 2D-coordinates
// representing the right, left, bottom, top neighbor (in this order, see TexelNeighbor::TexelNeighbor enum)
//
//
static void resolve_uv_neighborhood(const Surface_mesh& mesh, const cv::Mat& face_id_tex, const cv::Mat& halfedge_id_tex, Texel_neighborhood& resolved_neighborhood);

//
// Neighborhood access functions
//
static inline cv::Point top(const Texel_neighborhood& neighborhood, int row, int col);
static inline cv::Point left(const Texel_neighborhood& neighborhood, int row, int col);
static inline cv::Point right(const Texel_neighborhood& neighborhood, int row, int col);
static inline cv::Point bottom(const Texel_neighborhood& neighborhood, int row, int col);

//
// Validation functions
//
static inline bool valid(const Texel_neighborhood& neighborhood);
static inline bool valid(const Texel_neighborhood& neighborhood, int row, int col);


//
// IMPLEMENTATION
//

struct Texel_neighborhood
{
    std::vector<std::vector<cv::Point>> neighborhood;
    int width;
    int height;
};


// Helper enum for 4-neighborhood
namespace TexelNeighbor
{
    enum TexelNeighbor { R = 0, L = 1, B = 2, T = 3 };

    static TexelNeighbor opposite(TexelNeighbor t)
    {
        switch(t)
        {
            case R: return L;
            case L: return R;
            case B: return T;
            case T: return B;
        }

        return R;
    }
}

static inline int neighborhood_index(const Texel_neighborhood& neighborhood, int row, int col)
{
    int idx = row * neighborhood.width + col;
    return idx;
}

static inline cv::Point neighbor(const Texel_neighborhood& neighborhood, int row, int col, TexelNeighbor::TexelNeighbor neighbor)
{
    int idx = neighborhood_index(neighborhood, row, col);
    return neighborhood.neighborhood[idx][neighbor];
}

static inline cv::Point left(const Texel_neighborhood& neighborhood, int row, int col)
{
    return neighbor(neighborhood, row, col, TexelNeighbor::L);
}

static inline cv::Point right(const Texel_neighborhood& neighborhood, int row, int col)
{
    return neighbor(neighborhood, row, col, TexelNeighbor::R);
}

static inline cv::Point bottom(const Texel_neighborhood& neighborhood, int row, int col)
{
    return neighbor(neighborhood, row, col, TexelNeighbor::B);
}

static inline cv::Point top(const Texel_neighborhood& neighborhood, int row, int col)
{
    return neighbor(neighborhood, row, col, TexelNeighbor::T);
}

static inline bool valid(const Texel_neighborhood& neighborhood)
{
    return neighborhood.neighborhood.size() > 0;
}

static inline bool valid(const Texel_neighborhood& neighborhood, int row, int col)
{
    int idx = neighborhood_index(neighborhood, row, col);
    return neighborhood.neighborhood[idx].size() > 0;
}

//
// Looks through a 3x3 texel neighborhood and of all texels corresponding to valid texels (i.e. corresponding to a face_id)
// returns the texel closest to the given input
//
static bool round_to_next_valid_pixel(const double pix_i, const double pix_j, const cv::Mat &face_id_tex, graphene::Vec2ui &result)
{
    const int pix_i_ver1 = pix_i;
    const int pix_j_ver1 = pix_j;

    const int pix_i_ver2 = pix_i + 1;
    const int pix_j_ver2 = pix_j;

    const int pix_i_ver3 = pix_i;
    const int pix_j_ver3 = pix_j + 1;

    const int pix_i_ver4 = pix_i + 1;
    const int pix_j_ver4 = pix_j + 1;

    const int pix_i_ver5 = pix_i - 1;
    const int pix_j_ver5 = pix_j;

    const int pix_i_ver6 = pix_i;
    const int pix_j_ver6 = pix_j - 1;

    const int pix_i_ver7 = pix_i - 1;
    const int pix_j_ver7 = pix_j - 1;

    const int pix_i_ver8 = pix_i + 1;
    const int pix_j_ver8 = pix_j - 1;

    const int pix_i_ver9 = pix_i - 1;
    const int pix_j_ver9 = pix_j + 1;

    int faceID_ver1_pixel = face_id_tex.at<int32_t>(pix_j_ver1, pix_i_ver1);
    int faceID_ver2_pixel = face_id_tex.at<int32_t>(pix_j_ver2, pix_i_ver2);
    int faceID_ver3_pixel = face_id_tex.at<int32_t>(pix_j_ver3, pix_i_ver3);
    int faceID_ver4_pixel = face_id_tex.at<int32_t>(pix_j_ver4, pix_i_ver4);
    int faceID_ver5_pixel = face_id_tex.at<int32_t>(pix_j_ver5, pix_i_ver5);
    int faceID_ver6_pixel = face_id_tex.at<int32_t>(pix_j_ver6, pix_i_ver6);
    int faceID_ver7_pixel = face_id_tex.at<int32_t>(pix_j_ver7, pix_i_ver7);
    int faceID_ver8_pixel = face_id_tex.at<int32_t>(pix_j_ver8, pix_i_ver8);
    int faceID_ver9_pixel = face_id_tex.at<int32_t>(pix_j_ver9, pix_i_ver9);

    bool ver1_valid = faceID_ver1_pixel >= 0;
    bool ver2_valid = faceID_ver2_pixel >= 0;
    bool ver3_valid = faceID_ver3_pixel >= 0;
    bool ver4_valid = faceID_ver4_pixel >= 0;
    bool ver5_valid = faceID_ver5_pixel >= 0;
    bool ver6_valid = faceID_ver6_pixel >= 0;
    bool ver7_valid = faceID_ver7_pixel >= 0;
    bool ver8_valid = faceID_ver8_pixel >= 0;
    bool ver9_valid = faceID_ver9_pixel >= 0;

    const double diff_ver1 = std::pow(pix_i - pix_i_ver1, 2) + std::pow(pix_j - pix_j_ver1, 2);
    const double diff_ver2 = std::pow(pix_i - pix_i_ver2, 2) + std::pow(pix_j - pix_j_ver2, 2);
    const double diff_ver3 = std::pow(pix_i - pix_i_ver3, 2) + std::pow(pix_j - pix_j_ver3, 2);
    const double diff_ver4 = std::pow(pix_i - pix_i_ver4, 2) + std::pow(pix_j - pix_j_ver4, 2);
    const double diff_ver5 = std::pow(pix_i - pix_i_ver5, 2) + std::pow(pix_j - pix_j_ver5, 2);
    const double diff_ver6 = std::pow(pix_i - pix_i_ver6, 2) + std::pow(pix_j - pix_j_ver6, 2);
    const double diff_ver7 = std::pow(pix_i - pix_i_ver7, 2) + std::pow(pix_j - pix_j_ver7, 2);
    const double diff_ver8 = std::pow(pix_i - pix_i_ver8, 2) + std::pow(pix_j - pix_j_ver8, 2);
    const double diff_ver9 = std::pow(pix_i - pix_i_ver9, 2) + std::pow(pix_j - pix_j_ver9, 2);

    bool found_valid = false;
    graphene::Vec2ui best_pixel;

    if (ver1_valid || ver2_valid || ver3_valid || ver4_valid || ver5_valid ||
        ver6_valid || ver7_valid || ver8_valid || ver9_valid)
    {
        found_valid = true;

        double best_diff = DBL_MAX;
        if (ver1_valid && diff_ver1 < best_diff)
        {
            best_pixel = graphene::Vec2ui(pix_i_ver1, pix_j_ver1);
            best_diff = diff_ver1;
        }
        if (ver2_valid && diff_ver2 < best_diff)
        {
            best_pixel = graphene::Vec2ui(pix_i_ver2, pix_j_ver2);
            best_diff = diff_ver2;
        }
        if (ver3_valid && diff_ver3 < best_diff)
        {
            best_pixel = graphene::Vec2ui(pix_i_ver3, pix_j_ver3);
            best_diff = diff_ver4;
        }
        if (ver4_valid && diff_ver4 < best_diff)
        {
            best_pixel = graphene::Vec2ui(pix_i_ver4, pix_j_ver4);
            best_diff = diff_ver4;
        }
        if (ver5_valid && diff_ver5 < best_diff)
        {
            best_pixel = graphene::Vec2ui(pix_i_ver5, pix_j_ver5);
            best_diff = diff_ver5;
        }
        if (ver6_valid && diff_ver6 < best_diff)
        {
            best_pixel = graphene::Vec2ui(pix_i_ver6, pix_j_ver6);
            best_diff = diff_ver6;
        }
        if (ver7_valid && diff_ver7 < best_diff)
        {
            best_pixel = graphene::Vec2ui(pix_i_ver7, pix_j_ver7);
            best_diff = diff_ver7;
        }
        if (ver8_valid && diff_ver8 < best_diff)
        {
            best_pixel = graphene::Vec2ui(pix_i_ver8, pix_j_ver8);
            best_diff = diff_ver8;
        }
        if (ver9_valid && diff_ver9 < best_diff)
        {
            best_pixel = graphene::Vec2ui(pix_i_ver9, pix_j_ver9);
            best_diff = diff_ver9;
        }
    }
    else
    {
        found_valid = false;
        best_pixel = graphene::Vec2ui(pix_i_ver1, pix_j_ver1);
    }

    result = best_pixel;

    return found_valid;
}

//
// Resolves neighborhood on the uv layout of the passed mesh.
//
// First computes every texel whose neighborhood needs resolve (i.e. every texel whose standard 4-neighborhood contains invalid texels).
// For each of these, we compute the interpolation value on the corresponding uv-halfedge.
// Its neighbor is then computed by using this value to interpolate the opposite uv-halfedge.
//
static void resolve_uv_neighborhood(const Surface_mesh& mesh,
                                    const cv::Mat& face_id_tex,
                                    const cv::Mat& halfedge_id_tex,
                                    Texel_neighborhood& resolved_neighborhood)
{
    auto uvs = mesh.get_halfedge_property<graphene::Texture_coordinate>("h:texcoord");
    if (!uvs)
    {
        printf("[ERROR] Cannot get uv property for mesh\n");
        return;
    }

    const int width  = face_id_tex.cols;
    const int height = face_id_tex.rows;

    //
    // Setup neighborhood struct
    //
    resolved_neighborhood.width = width;
    resolved_neighborhood.height = height;
    resolved_neighborhood.neighborhood.resize(width * height);

    //
    // Initialize with standard 4 neighborhood and remember texels that need resolve
    //
    cv::Mat needs_neighborhood_resolve = cv::Mat::zeros(width, height, CV_8UC4);

    for (int row = 1; row < height - 1; ++row)
    {
        for (int col = 1; col < width - 1; ++col)
        {
            int id = face_id_tex.at<int32_t>(row, col);

            // Every valid texel has 4 neighbors
            int index = neighborhood_index(resolved_neighborhood, row, col);
            resolved_neighborhood.neighborhood[index].resize(4);

            // Initialize neighborhood with standard 4 neighborhood
            resolved_neighborhood.neighborhood[index][TexelNeighbor::R] = cv::Point(col + 1, row    );
            resolved_neighborhood.neighborhood[index][TexelNeighbor::L] = cv::Point(col - 1, row    );
            resolved_neighborhood.neighborhood[index][TexelNeighbor::B] = cv::Point(col    , row + 1);
            resolved_neighborhood.neighborhood[index][TexelNeighbor::T] = cv::Point(col    , row - 1);

            //
            // Remember which neighbors need resolve
            // Every texel with a corresponding face id and whose neighbors
            // do not have a face id need resolve
            //
            if (id > -1)
            {
                auto& needs_resolve = needs_neighborhood_resolve.at<cv::Vec4b>(row, col);

                if (face_id_tex.at<int32_t>(row    , col + 1) < 0) needs_resolve[TexelNeighbor::R] = 255;
                if (face_id_tex.at<int32_t>(row    , col - 1) < 0) needs_resolve[TexelNeighbor::L] = 255;
                if (face_id_tex.at<int32_t>(row + 1, col    ) < 0) needs_resolve[TexelNeighbor::B] = 255;
                if (face_id_tex.at<int32_t>(row - 1, col    ) < 0) needs_resolve[TexelNeighbor::T] = 255;
            }
        }
    }

    //
    // Now that we know which texels need their neighborhood resolve,
    // actually compute the resolved neighbors
    //
    for (int row = 1; row < height - 1; ++row)
    {
        for (int col = 1; col < width - 1; ++col)
        {
            int index = neighborhood_index(resolved_neighborhood, row, col);
            cv::Vec4b neighbors_to_resolve = needs_neighborhood_resolve.at<cv::Vec4b>(row, col);

            for (int n = 0; n < 4; ++n)
            {
                bool needs_resolve = neighbors_to_resolve[n] > 0;
                if (!needs_resolve)
                    continue;

                // Which halfedge belongs to this texel
                int32_t halfedge_id = halfedge_id_tex.at<int32_t>(row, col);
                Surface_mesh::Halfedge he(halfedge_id);
                Surface_mesh::Halfedge opposite_he = mesh.opposite_halfedge(he);

                if (mesh.is_boundary(opposite_he))
                {
                    resolved_neighborhood.neighborhood[index][n] = cv::Point(col, row);
                    continue;
                }

                // Top-down to bottom-up, pixel to uv
                const graphene::Texture_coordinate uv = {(float)col / ((float)width),
                                                         1.0f - ((float)row / ((float)height)),
                                                         1.0f};

                // Corresponding edge in uv space
                auto uv1 = uvs[he];
                auto uv2 = uvs[mesh.prev_halfedge(he)];

                // Opposite edge in uv space
                auto uv3 = uvs[opposite_he];
                auto uv4 = uvs[mesh.prev_halfedge(opposite_he)];

                // Compute interpolation value of current uv on uv edge
                float t = graphene::distance(uv1, uv) / graphene::distance(uv1, uv2);
                // clamp between 0 and 1
                t = std::max(std::min(t, 1.0f), 0.0f);

                // Interpolate on opposite edge
                auto neighbor_uv = uv3 + (1.0f - t) * (uv4 - uv3);

                // Bottom-up to top-down, uv to pixel
                neighbor_uv[1] = 1.0f - neighbor_uv[1];
                neighbor_uv[0] *= width;
                neighbor_uv[1] *= height;

                // Find the closest valid integer pixel coordinate to the computed floating point pixel coordinate
                graphene::Vec2ui integer_texel;
                if (!round_to_next_valid_pixel(neighbor_uv[0], neighbor_uv[1], face_id_tex, integer_texel))
                {
                    //std::cerr << "[ERROR] CANNOT FIND VALID PIXEL! SHOULD NOT HAPPEN !!!" << std::endl;

                    // Use standard rounding instead of explicit rounding
                    resolved_neighborhood.neighborhood[index][n] = cv::Point(neighbor_uv[0], neighbor_uv[1]);
                }
                else
                {
                    // Use the found integer coordinate
                    resolved_neighborhood.neighborhood[index][n] = cv::Point(integer_texel[0], integer_texel[1]);
                }
            }
        }
    }
}

#if 0

//
// Function, that goes through 3D to resolve uv neighborhood
//


static Point surface_position_of_texel(const Surface_mesh& mesh,
                                       const Surface_mesh::Face& face,
                                       const Surface_mesh::Halfedge_property<graphene::Texture_coordinate>& uvs,
                                       const int row,
                                       const int col,
                                       const int width,
                                       const int height)
{
    auto hiter = mesh.halfedges(face);
    const auto h0 = *hiter;
    const auto h1 = *(++hiter);
    const auto h2 = *(++hiter);
    const auto uv0 = uvs[h0];
    const auto uv1 = uvs[h1];
    const auto uv2 = uvs[h2];
    const auto p0 = Point(uv0[0], uv0[1], 1.0f);
    const auto p1 = Point(uv1[0], uv1[1], 1.0f);
    const auto p2 = Point(uv2[0], uv2[1], 1.0f);

    const Point uv = {(float)col / ((float)width), 1.0f - ((float)row / ((float)height)), 1.0f};
    // Barycentric coordinates of this texel in the corresponding uv triangle
    const Point b = graphene::geometry::barycentric_coordinates(uv, p0, p1, p2);

    const auto p0_mesh = mesh.position(mesh.to_vertex(h0));
    const auto p1_mesh = mesh.position(mesh.to_vertex(h1));
    const auto p2_mesh = mesh.position(mesh.to_vertex(h2));

    // Use the barycentric coordinates of the uv layout to interpolate the 3D position
    Point p = b[0] * p0_mesh + b[1] * p1_mesh + b[2] * p2_mesh;

    return p;
}


static void resolve_uv_patch_neighborhood(const Surface_mesh& mesh,
                                          const cv::Mat& face_id_tex,
                                          std::vector<std::vector<cv::Point>>& resolved_neighborhood)
{
    const int width  = face_id_tex.cols;
    const int height = face_id_tex.rows;

    // NOTE(swenninger): wastes memory by also taking invalid texels into account
    resolved_neighborhood.resize(width * height);

    // Each texel, whose neighbor does not belong to a valid face id needs its neighborhood resolved
    cv::Mat needs_neighborhood_resolve = cv::Mat::zeros(width, height, CV_8UC4);

    auto uvs = mesh.get_halfedge_property<graphene::Texture_coordinate>("h:texcoord");
    if (!uvs)
    {
        printf("[ERROR] Cannot get uv property for mesh\n");
        return;
    }

    for (int row = 1; row < height - 1; ++row)
    {
        for (int col = 1; col < width - 1; ++col)
        {
            int id = decode_face_id(face_id_tex.at<int32_t>(row, col));

            if (id < 0) // texel does not belong to any face
                continue;

            // Every valid texel has 4 neighbors
            int neighborhood_index = row * width + col;
            resolved_neighborhood[neighborhood_index].resize(4);
            // Initialize neighborhood with standard 4 neighborhood
            resolved_neighborhood[neighborhood_index][TexelNeighbor::R] = cv::Point(col + 1, row    );
            resolved_neighborhood[neighborhood_index][TexelNeighbor::L] = cv::Point(col - 1, row    );
            resolved_neighborhood[neighborhood_index][TexelNeighbor::B] = cv::Point(col    , row + 1);
            resolved_neighborhood[neighborhood_index][TexelNeighbor::T] = cv::Point(col    , row - 1);

            // Remember which neighbors need resolve
            auto& needs_resolve = needs_neighborhood_resolve.at<cv::Vec4b>(row, col);

            if (decode_face_id(face_id_tex.at<int32_t>(row    , col + 1)) < 0) needs_resolve[TexelNeighbor::R] = 255;
            if (decode_face_id(face_id_tex.at<int32_t>(row    , col - 1)) < 0) needs_resolve[TexelNeighbor::L] = 255;
            if (decode_face_id(face_id_tex.at<int32_t>(row + 1, col    )) < 0) needs_resolve[TexelNeighbor::B] = 255;
            if (decode_face_id(face_id_tex.at<int32_t>(row - 1, col    )) < 0) needs_resolve[TexelNeighbor::T] = 255;
        }
    }

    // Now that we know which texels need their neighborhood resolve, actually compute the neighbors
    for (int row = 1; row < height - 1; ++row)
    {
        for (int col = 1; col < width - 1; ++col)
        {
            cv::Vec4b neighbors_to_resolve = needs_neighborhood_resolve.at<cv::Vec4b>(row, col);

            for (int n = 0; n < 4; ++n)
            {
                bool needs_resolve = neighbors_to_resolve[n] > 0;

                if (!needs_resolve) continue;

                // Which face belongs to this texel
                int id = decode_face_id(face_id_tex.at<int32_t>(row, col));
                Surface_mesh::Face face_handle = Surface_mesh::Face(id);

                // Current 3D position corresponding to this texel
                Point p = surface_position_of_texel(mesh, face_handle, uvs, row, col, width, height);

                // Neighboring triangles of face. Check these for closest point to current 3D position
                std::set<Surface_mesh::Face> neighbors;

                for (auto v : mesh.vertices(face_handle))
                    for (auto f : mesh.faces(v))
                        neighbors.insert(f);

                neighbors.erase(face_handle);

                // Find the closest point to the current 3D position in all neighboring faces
                // This closest point defines the new neighbor
                double min_dist = DBL_MAX;
                Point closest_point;                 // Closest point on neighboring faces
                Surface_mesh::Face closest_neighbor; // Face belonging to closest point
                Point b;                             // barycentric coordinates of closest point in closest face
                for (auto f : neighbors)
                {
                    auto viter = mesh.vertices(f);
                    auto v0 = *viter;
                    auto v1 = *(++viter);
                    auto v2 = *(++viter);

                    const auto p0 = mesh.position(v0);
                    const auto p1 = mesh.position(v1);
                    const auto p2 = mesh.position(v2);

                    Point candidate;
                    double dist = graphene::geometry::dist_point_triangle(p, p0, p1, p2, candidate);

                    if (min_dist > dist)
                    {
                        min_dist = dist;
                        closest_neighbor = f;
                        closest_point = candidate;
                        b = graphene::geometry::barycentric_coordinates(closest_point, p0, p1, p2);
                    }
                }

                // For some texels, the numerical computation returns barycentric coordinates only slightly negative
                if (b[0] < -0.0001 || b[0] > 1.0001 || b[1] < -0.0001 || b[1] > 1.0001 || b[2] < -0.0001 || b[2] > 1.0001)
                {
                    std::cerr << "[ERROR] SHOULD NOT HAPPEN !!!" << std::endl;
                    std::cerr << "b   " << b << std::endl;
                    std::cerr << "p   " << p << std::endl;
                    std::cerr << "id  " << id << std::endl;
                    std::cerr << "c_n " << closest_neighbor << std::endl;
                    std::cerr << "c_p " << closest_point << std::endl;
                   // continue;
                }

                // Clamp barycentric coordinates
                if (b[0] < 0.0) b[0] = 0.0;
                if (b[0] > 1.0) b[0] = 1.0;
                if (b[1] < 0.0) b[1] = 0.0;
                if (b[1] > 1.0) b[1] = 1.0;
                if (b[2] < 0.0) b[2] = 0.0;
                if (b[2] > 1.0) b[2] = 1.0;

                // Compute uv coordinates of closest point via barycentric interpolation
                graphene::Texture_coordinate c(0.0f, 0.0f, 0.0f);
                int index = 0;
                for (auto he : mesh.halfedges(closest_neighbor))
                {
                    auto uv = uvs[he];
                    c += b[index++] * uv;
                }

                // flip y
                c[1] = 1.0f - c[1];

                // uv coordinates to (floating point) texel coordinates
                c[0] *= width;
                c[1] *= height;

                int neighborhood_index = row * width + col;

                // Find the closest integer pixel coordinate to the computed floating point pixel coordinate
                graphene::Vec2ui integer_texel;
                if (!round_to_next_valid_pixel(c[0], c[1], face_id_tex, integer_texel))
                {
                    std::cerr << "[ERROR] CANNOT FIND VALID PIXEL! SHOULD NOT HAPPEN !!!" << std::endl;
                    resolved_neighborhood[neighborhood_index][n] = cv::Point(0,0);
                    continue;
                }
                else
                {
                    resolved_neighborhood[neighborhood_index][n] = cv::Point(integer_texel[0], integer_texel[1]);
                }
            }
        }
    }
}
#endif
