#include <cnpy.h>

#include <CLI/CLI.hpp>

#include "../classes.h"
#include "../functions.h"

extern pangolin::GlSlProgram GetShaderProgram();

void write_to_npz(const int index,
                  const pangolin::Image<Eigen::Vector4f> &image,
                  const std::string path) {
    std::vector<float> data;
    for (unsigned int w = 0; w < image.w; w++) {
        for (unsigned int h = 0; h < image.h; h++) {
            Eigen::Vector4f normal = image(w, h);
            data.push_back(normal[0]);
            data.push_back(normal[1]);
            data.push_back(normal[2]);
        }
    }
    std::string key = "view_" + std::to_string(index + 1);
    if (index == 0) {
        cnpy::npz_save(path, key, &data[0],
                       {(long unsigned int)(data.size() / 3), 3}, "w");
    } else {
        cnpy::npz_save(path, key, &data[0],
                       {(long unsigned int)(data.size() / 3), 3}, "a");
    }
}

int main(int argc, char **argv) {
    std::string mesh_path;
    std::string npz_path;
    int num_viewpoints = 100;
    float radius = 1;
    size_t width = 400;
    size_t height = 400;

    CLI::App app{"surface_renderer"};
    app.add_option("--input-mesh-path", mesh_path)->required();
    app.add_option("--output-npz-path", npz_path)->required();
    app.add_option("--num-viewpoints", num_viewpoints);
    app.add_option("--image-width", width);
    app.add_option("--image-height", height);
    CLI11_PARSE(app, argc, argv);

    glPixelStorei(GL_UNPACK_ALIGNMENT, 1);
    glPixelStorei(GL_UNPACK_ROW_LENGTH, 0);
    glPixelStorei(GL_UNPACK_SKIP_PIXELS, 0);
    glPixelStorei(GL_UNPACK_SKIP_ROWS, 0);

    pangolin::Geometry geometry = pangolin::LoadGeometry(mesh_path);
    std::cout << geometry.objects.size() << " objects" << std::endl;
    geometry.textures.clear();

    int total_num_faces = 0;
    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);

            total_num_faces += ibo.h;
        }
    }
    std::cout << total_num_faces << " faces" << std::endl;

    pangolin::ManagedImage<uint8_t> new_buffer(3 * sizeof(uint32_t),
                                               total_num_faces);
    pangolin::Image<uint32_t> new_ibo =
        new_buffer.UnsafeReinterpret<uint32_t>().SubImage(0, 0, 3,
                                                          total_num_faces);
    int new_index = 0;
    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 k = 0; k < ibo.h; ++k) {
                new_ibo.Row(new_index).CopyFrom(ibo.Row(k));
                new_index++;
            }
        }
    }
    geometry.objects.clear();

    auto faces = geometry.objects.emplace(std::string("mesh"),
                                          pangolin::Geometry::Element());
    faces->second.Reinitialise(3 * sizeof(uint32_t), total_num_faces);
    faces->second.CopyFrom(new_buffer);
    new_ibo = faces->second.UnsafeReinterpret<uint32_t>().SubImage(
        0, 0, 3, total_num_faces);
    faces->second.attributes["vertex_indices"] = new_ibo;

    auto params = normalize_geometry(geometry, true);

    pangolin::CreateWindowAndBind("Main", 1, 1);
    glEnable(GL_DEPTH_TEST);
    glDisable(GL_DITHER);
    glDisable(GL_POINT_SMOOTH);
    glDisable(GL_LINE_SMOOTH);
    glDisable(GL_POLYGON_SMOOTH);
    glHint(GL_POINT_SMOOTH, GL_DONT_CARE);
    glHint(GL_LINE_SMOOTH, GL_DONT_CARE);
    glHint(GL_POLYGON_SMOOTH_HINT, GL_DONT_CARE);
    glDisable(GL_MULTISAMPLE_ARB);
    glShadeModel(GL_FLAT);

    // Define Projection and initial ModelView matrix
    float camera_distance = 1;
    pangolin::OpenGlRenderState camera(
        pangolin::ProjectionMatrixOrthographic(-camera_distance,
                                               camera_distance, camera_distance,
                                               -camera_distance, 0, 2.5),
        pangolin::ModelViewLookAt(0, 0, -1, 0, 0, 0, pangolin::AxisY));
    pangolin::GlGeometry gl_geometry = pangolin::ToGlGeometry(geometry);
    pangolin::GlSlProgram program = GetShaderProgram();

    pangolin::GlRenderBuffer z_buffer(width, height, GL_DEPTH_COMPONENT32);
    pangolin::GlTexture texture_normals(width, height, GL_RGBA32F);
    pangolin::GlTexture texture_vertices(width, height, GL_RGBA32F);
    pangolin::GlFramebuffer frame_buffer(texture_normals, texture_vertices,
                                         z_buffer);

    SurfaceRenderer renderer(geometry, width, height, camera_distance);
    SurfaceRendererState state(total_num_faces, num_viewpoints,
                               camera_distance * 1.1);

    for (unsigned int view_index = 0; view_index < state.views.size();
         view_index++) {
        auto view = state.views[view_index];
        pangolin::TypedImage image_normals;
        pangolin::TypedImage image_vertices;
        renderer.render_from_view(
            view_index, view, program, texture_normals, texture_vertices,
            frame_buffer, image_normals, image_vertices, state.observed_normals,
            state.observed_points, state.observed_points_view_index,
            state.face_normal_test, state.invisible_face_test,
            state.observed_backface_test, state.face_observation_count);
        write_to_npz(view_index,
                     image_normals.UnsafeReinterpret<Eigen::Vector4f>(),
                     npz_path);
    }

    return 0;
}
