#include "functions.h"

#include <random>

std::vector<Eigen::Vector3f> sample_points_on_unit_sphere(
    const uint num_samples, const float radius) {
    std::vector<Eigen::Vector3f> points(num_samples);
    const float offset = 2.0f / num_samples;
    const float increment = static_cast<float>(M_PI) * (3.0f - std::sqrt(5.0f));

    for (uint k = 0; k < num_samples; k++) {
        const float y = ((k * offset) - 1) + (offset / 2);
        const float r = std::sqrt(1 - std::pow(y, 2.0f));
        const float phi = (k + 1.0f) * increment;
        const float x = cos(phi) * r;
        const float z = sin(phi) * r;
        points[k] = radius * Eigen::Vector3f(x, y, z);
    }
    return points;
}

void find_invisible_faces(const pangolin::Image<Eigen::Vector4f>& image_normals,
                          std::vector<bool>& invisible_face_test) {
    Eigen::Vector4f normal;
    for (unsigned int w = 0; w < image_normals.w; w++) {
        for (unsigned int h = 0; h < image_normals.h; h++) {
            normal = image_normals(w, h);
            if (normal[3] == 0.0f) {
                continue;
            }
            const auto face_index =
                static_cast<std::size_t>(normal[3] + 0.01f) - 1;
            invisible_face_test[face_index] = false;
        }
    }
}

void find_observed_backfaces(
    const pangolin::Image<Eigen::Vector4f>& image_normals,
    std::vector<Eigen::Vector4f>& face_normal_test,
    std::vector<bool>& observed_backface_test,
    std::vector<int>& face_observation_count) {
    Eigen::Vector4f rendered_normal;
    for (unsigned int w = 0; w < image_normals.w; w++) {
        for (unsigned int h = 0; h < image_normals.h; h++) {
            rendered_normal = image_normals(w, h);
            if (rendered_normal[3] == 0.0f) {
                continue;
            }
            const auto face_index =
                static_cast<std::size_t>(rendered_normal[3] + 0.01f) - 1;
            face_observation_count[face_index] += 1;
            Eigen::Vector4f face_normal = face_normal_test[face_index];
            if (face_normal[3] == 0.0f) {
                face_normal_test[face_index] = rendered_normal;  // init
            } else if (face_normal[3] > 0.0f) {
                const float dot =
                    face_normal.head<3>().dot(rendered_normal.head<3>());
                if (dot < 0.0f) {
                    face_normal_test[face_index][3] = -1.0f;
                    observed_backface_test[face_index] = true;
                }
            }
        }
    }
}

void get_valid_surface_points_and_normals_from_image(
    const int view_index, const pangolin::Image<Eigen::Vector4f>& image_normals,
    const pangolin::Image<Eigen::Vector4f>& image_vertices,
    std::vector<Eigen::Vector4f>& surface_normals,
    std::vector<Eigen::Vector4f>& surface_points,
    std::vector<int>& surface_points_view_index) {
    Eigen::Vector4f vertex;
    Eigen::Vector4f normal;
    for (unsigned int w = 0; w < image_vertices.w; w++) {
        for (unsigned int h = 0; h < image_vertices.h; h++) {
            normal = image_normals(w, h);
            if (normal[3] == 0.0f) {
                continue;
            }
            vertex = image_vertices(w, h);
            if (vertex[3] == 0.0f) {
                continue;
            }
            surface_normals.push_back(normal);
            surface_points.push_back(vertex);
            surface_points_view_index.push_back(view_index);
        }
    }
}

float compute_triangle_area(const Eigen::Vector3f& a, const Eigen::Vector3f& b,
                            const Eigen::Vector3f& c) {
    const Eigen::Vector3f ab = b - a;
    const Eigen::Vector3f ac = c - a;

    float costheta = ab.dot(ac) / (ab.norm() * ac.norm());

    if (costheta < -1)  // meaning theta is pi
        costheta = std::cos(static_cast<float>(M_PI) * 359.f / 360);
    else if (costheta > 1)  // meaning theta is zero
        costheta = std::cos(static_cast<float>(M_PI) * 1.f / 360);

    const float sinTheta = std::sqrt(1 - costheta * costheta);

    return 0.5f * ab.norm() * ac.norm() * sinTheta;
}

Eigen::Vector3f sample_point_from_triangle(const Eigen::Vector3f& a,
                                           const Eigen::Vector3f& b,
                                           const Eigen::Vector3f& c) {
    std::random_device seeder;
    std::mt19937 generator(seeder());
    std::uniform_real_distribution<float> rand_dist(0.0, 1.0);

    const float r1 = rand_dist(generator);
    const float r2 = rand_dist(generator);

    return Eigen::Vector3f((1 - std::sqrt(r1)) * a +
                           std::sqrt(r1) * (1 - r2) * b +
                           r2 * std::sqrt(r1) * c);
}

GeometryNormalizationParameters normalize_geometry(pangolin::Geometry& geometry,
                                                   bool fit_to_unit_sphere,
                                                   const float buffer) {
    float x_min = 1000000, x_max = -1000000, y_min = 1000000, y_max = -1000000,
          z_min = 1000000, z_max = -1000000;

    pangolin::Image<float> vertices = pangolin::get<pangolin::Image<float>>(
        geometry.buffers["geometry"].attributes["vertex"]);

    const std::size_t num_vertices = vertices.h;

    std::vector<bool> used_vertices(num_vertices, false);
    for (const auto& object : geometry.objects) {
        auto vertex_indices_iter =
            object.second.attributes.find("vertex_indices");
        if (vertex_indices_iter != object.second.attributes.end()) {
            pangolin::Image<uint32_t> ibo =
                pangolin::get<pangolin::Image<uint32_t>>(
                    vertex_indices_iter->second);
            for (uint face_index = 0; face_index < ibo.h; face_index++) {
                for (uint n = 0; n < 3; n++) {
                    used_vertices[ibo(n, face_index)] = true;
                }
            }
        }
    }

    for (size_t k = 0; k < num_vertices; k++) {
        if (used_vertices[k] == false) {
            continue;
        }
        x_min = fmin(x_min, vertices(0, k));
        y_min = fmin(y_min, vertices(1, k));
        z_min = fmin(z_min, vertices(2, k));
        x_max = fmax(x_max, vertices(0, k));
        y_max = fmax(y_max, vertices(1, k));
        z_max = fmax(z_max, vertices(2, k));
    }
    const float x_center = (x_max + x_min) / 2.0f;
    const float y_center = (y_max + y_min) / 2.0f;
    const float z_center = (z_max + z_min) / 2.0f;

    float max_distance = -1.0f;
    for (size_t k = 0; k < num_vertices; k++) {
        if (used_vertices[k] == false) {
            continue;
        }
        vertices(0, k) -= x_center;
        vertices(1, k) -= y_center;
        vertices(2, k) -= z_center;
        const float distance =
            Eigen::Map<Eigen::Vector3f>(vertices.RowPtr(k)).norm();
        max_distance = std::max(max_distance, distance);
    }

    max_distance *= buffer;

    if (fit_to_unit_sphere) {
        for (size_t k = 0; k < num_vertices; k++) {
            vertices(0, k) /= max_distance;
            vertices(1, k) /= max_distance;
            vertices(2, k) /= max_distance;
        }
    }
    GeometryNormalizationParameters params;
    params.center = {x_center, y_center, z_center};
    params.max_distance = max_distance;
    return params;
}

std::vector<Eigen::Vector3i> get_geometry_faces(pangolin::Geometry& geometry) {
    std::vector<Eigen::Vector3i> faces;
    for (const auto& object : geometry.objects) {
        auto vertext_indices = object.second.attributes.find("vertex_indices");
        if (vertext_indices != object.second.attributes.end()) {
            pangolin::Image<uint32_t> ibo =
                pangolin::get<pangolin::Image<uint32_t>>(
                    vertext_indices->second);
            for (int i = 0; i < ibo.h; ++i) {
                faces.emplace_back(ibo(0, i), ibo(1, i), ibo(2, i));
            }
        }
    }
    return faces;
}

std::vector<float> compute_face_area(pangolin::Geometry& geometry,
                                     std::vector<bool>& ignore_face) {
    float total_area = 0.0f;

    std::vector<Eigen::Vector3i> faces = get_geometry_faces(geometry);
    pangolin::Image<float> vertices = pangolin::get<pangolin::Image<float>>(
        geometry.buffers["geometry"].attributes["vertex"]);

    std::vector<float> face_area(faces.size());
    for (unsigned int face_index = 0; face_index < faces.size(); face_index++) {
        float area = 0;
        if (ignore_face[face_index] == false) {
            const Eigen::Vector3i& face = faces[face_index];
            area = compute_triangle_area(
                (Eigen::Vector3f)Eigen::Map<Eigen::Vector3f>(
                    vertices.RowPtr(face(0))),
                (Eigen::Vector3f)Eigen::Map<Eigen::Vector3f>(
                    vertices.RowPtr(face(1))),
                (Eigen::Vector3f)Eigen::Map<Eigen::Vector3f>(
                    vertices.RowPtr(face(2))));
        }
        face_area[face_index] = area;
    }
    return face_area;
}

void sample_from_surface(pangolin::Geometry& geometry,
                         std::vector<bool>& ignore_face,
                         std::vector<Eigen::Vector3f>& face_normals,
                         std::vector<Eigen::Vector3f>& sampled_normals,
                         std::vector<Eigen::Vector3f>& sampled_points,
                         int num_sample) {
    if (geometry.objects.size() != 1) {
        std::cout << "geometry.objects.size() != 1" << std::endl;
        return;
    }

    std::vector<Eigen::Vector3i> faces = get_geometry_faces(geometry);
    pangolin::Image<float> vertices = pangolin::get<pangolin::Image<float>>(
        geometry.buffers["geometry"].attributes["vertex"]);

    if (faces.size() != ignore_face.size()) {
        std::cout << "faces.size() != backface_test.size()" << std::endl;
        return;
    }

    float total_area = 0.0f;
    std::vector<float> face_area = compute_face_area(geometry, ignore_face);
    std::vector<float> cdf_by_area;

    for (unsigned int face_index = 0; face_index < faces.size(); face_index++) {
        float area = 0;
        if (ignore_face[face_index] == false) {
            area = face_area[face_index];
        }
        if (std::isnan(area)) {
            area = 0;
        }
        total_area += area;

        if (cdf_by_area.empty()) {
            cdf_by_area.push_back(area);
        } else {
            cdf_by_area.push_back(cdf_by_area.back() + area);
        }
    }

    std::cout << "total_area: " << total_area << std::endl;

    std::random_device seeder;
    std::mt19937 generator(seeder());
    std::uniform_real_distribution<float> rand_dist(0.0, total_area);

    while ((int)sampled_points.size() < num_sample) {
        float face_sample = rand_dist(generator);
        std::vector<float>::iterator face_index_iter = std::lower_bound(
            cdf_by_area.begin(), cdf_by_area.end(), face_sample);
        int face_index = std::distance(cdf_by_area.begin(), face_index_iter);
        if (ignore_face[face_index] == true) {
            continue;
        }

        const Eigen::Vector3i& face = faces[face_index];

        const Eigen::Vector3f va =
            Eigen::Map<Eigen::Vector3f>(vertices.RowPtr(face(0)));
        const Eigen::Vector3f vb =
            Eigen::Map<Eigen::Vector3f>(vertices.RowPtr(face(1)));
        const Eigen::Vector3f vc =
            Eigen::Map<Eigen::Vector3f>(vertices.RowPtr(face(2)));

        sampled_points.push_back(sample_point_from_triangle(va, vb, vc));
        sampled_normals.push_back(face_normals[face_index]);
    }
}
