#include "ofApp.h"

//--------------------------------------------------------------
void ofApp::setup(){

    /* allocated fbo's */
    map_fbo.allocate(ofGetWindowWidth() / 2, map_h, GL_RGB);
    map_fbo_alpha.allocate(map_fbo.getWidth(), map_h, GL_RGBA);
    map_fbo_post.allocate(map_fbo.getWidth(), map_h, GL_RGBA);

    portrait_fbo.allocate((ofGetWindowWidth() / 2), map_h, GL_RGBA); // this shoould be full with
    portrait_pre_fbo.allocate((ofGetWindowWidth() / 3) * 1, map_h, GL_RGB); // this should be full, and same below
    portrait_pre_fbo_alpha.allocate((ofGetWindowWidth() / 3) * 1, map_h, GL_RGBA);
    portrait_cropped.allocate((ofGetWindowWidth() / 3) * 1, map_h, GL_RGB);
    portrait_cropped_alpha.allocate((ofGetWindowWidth() / 3) * 1, map_h, GL_RGBA);

    comp_fbo.allocate(ofGetWindowWidth(), map_h, GL_RGBA);

    model_outptut_fbo.allocate(map_fbo.getWidth(), map_h, GL_RGB);
    model_esp_out_fbo.allocate(128 * 2, 168 * 2, GL_RGB);
    model_portrait_out_fbo.allocate(portrait_pre_fbo_alpha.getWidth(), portrait_pre_fbo_alpha.getHeight(), GL_RGB);

    /* k-d image comp (4-images) */
    kd_out.allocate(128 * 2, 168 * 2, GL_RGB);
    esp_comp_fbo.allocate(128 * 2, 168 * 2, GL_RGB);

    /* input images for model */
    model_image_esp.allocate(128 * 2, 168 * 2, OF_IMAGE_COLOR);
    model_image.allocate(map_fbo.getWidth(), map_h, OF_IMAGE_COLOR);
    model_image_portrait.allocate(portrait_pre_fbo_alpha.getWidth(), portrait_pre_fbo_alpha.getHeight(), OF_IMAGE_COLOR);
    model_image_portrait_cropped.allocate(portrait_pre_fbo_alpha.getWidth(), portrait_pre_fbo_alpha.getHeight(), OF_IMAGE_COLOR);

    alpha_demo.allocate(portrait_pre_fbo_alpha.getWidth(), portrait_pre_fbo_alpha.getHeight(), OF_IMAGE_COLOR_ALPHA);

    /* allocated ofImages in esp_images */
    for(int i = 0; i < 4; i++){
        ofFbo temp;
        temp.allocate(128, 168, GL_RGB);
        esp_images.push_back(temp);
    }

    /* of settings (not sure why putting this at the start of setup breaks the map - but it works!) */
    
    ofSetVerticalSync(true);
    ofDisableArbTex();
    ofEnableDepthTest();
    ofDisableAntiAliasing();
    ofEnableAlphaBlending();

    /* load */
    ofLoadImage(bayer, "images/bayer.png");
    bayer.setTextureWrap(GL_REPEAT, GL_REPEAT);
    shaders.load("shaders/dither");
    esp_shader.load("shaders/espDepth");
    vert_snap.load("shaders/vertex_snap");
    p_depth.load("shaders/p_depth");
    map_depth.load("shaders/map_depth");
   
    ORTCHAR_T* modelPath = "/home/cailean/Desktop/openframeworks/of_v0.12.0_linux64gcc6_release/apps/myApps/image-to-mesh/bin/data/models/depth_anything_v2_vits.onnx";
    ORTCHAR_T* modelPath_Small = "/home/cailean/Desktop/openframeworks/of_v0.12.0_linux64gcc6_release/apps/myApps/image-to-mesh/bin/data/models/depth_anything_v2_vitb.onnx";
    ORTCHAR_T* modelPath_Yolo = "/home/cailean/Desktop/openframeworks/of_v0.12.0_linux64gcc6_release/apps/myApps/image-to-mesh/bin/data/models/yolov5s-face.onnx";
    ORTCHAR_T* modelPath_Indoor_Dynamic = "/home/cailean/Desktop/openframeworks/of_v0.12.0_linux64gcc6_release/apps/myApps/image-to-mesh/bin/data/models/depth_anything_v2_vitl.onnx";

    bullet.setup();

    /* pi mapper setup */
    // mapper.registerFboSource(projection_src);
    // mapper.setup();
    // projection_src.setFbo(&comp_fbo);

    /* onnx setup */
    depth_onnx.Setup(modelPath, false, true);
    depth_onnx_esp.Setup(modelPath, false, true);
    depth_onnx_portrait.Setup(modelPath, false, true);
    yolo_onnx.Setup(modelPath_Yolo, false, true);

    /* multi-thread setup */
    depth_thread.setup(&model_image, &model_outptut_fbo, &depth_onnx);
    depth_esp.setup(&model_image_esp, &model_esp_out_fbo, &depth_onnx_esp);
    depth_portrait.setup(&model_image_portrait_cropped, &model_portrait_out_fbo, &depth_onnx_portrait);
    yolo.setupYolo(&model_image_portrait, &detected_faces, &yolo_onnx, &face_detector);
    
    /* camera settings for portrait */
    CameraPosition cp;
    cp.position = glm::vec3(-131.228, 502.624, 706.1);
    cp.rotation = glm::quat(0.898726, -0.436474, -0.0351978, -0.0170607);
    cp.scale = 0.313542;
    cam_positions.push_back(cp);

    cp.position = glm::vec3(7.987, -195.329, 813.56);
    cp.rotation = glm::quat(0.993, 0.117, 0.005, -0.00057);
    cam_positions.push_back(cp);

    cp.position = glm::vec3(160.81, -87.86, 731.575);
    cp.rotation = glm::quat(0.980, 0.027, 0.1934, -0.0053);
    cam_positions.push_back(cp);

    /* settings */
    portrait_camera.enableOrtho();
    portrait_camera.setNearClip(-10000);
    portrait_camera.setFarClip(10000);
    portrait_camera.setPosition(cam_positions[0].position);
    portrait_camera.setOrientation(cam_positions[0].rotation);
    portrait_camera.setScale(cam_positions[0].scale);
    portrait_camera.removeAllInteractions();
    portrait_camera.disableMouseInput();
    
    createNodes("data/json/sv_embeddings.json");

    buildMeshes();

    buildKDTree();

    bullet.setNodes(nodes);

    server = std::make_unique<Server>(6762, embed, vp_tree, nodes, false, "192.168.0.253", 2000, "search");
    server->start();

    last_updated_time = ofGetElapsedTimef();
    last_updated_time_esp = ofGetElapsedTimef();
}

//--------------------------------------------------------------
void ofApp::update(){
    server->update(esp_comp_fbo);

    float current_time = ofGetElapsedTimef();

    Node n = server->getChosenNode();

    // Check if node has changed
    if(n.tsne_position != last_chosen_node.tsne_position) {
        portrait_needs_update = true;
        last_chosen_node = n;
    } else {
        portrait_needs_update = false;
    }


    if(current_time - last_updated_time >= 0.1){
        buildKDTree();
        getNearestImages();
        last_updated_time = current_time;
    }
    
    if(current_time - last_updated_time_esp >= 4.0){
        depth_esp.update();
        depth_onnx_esp.SetPixels(model_esp_out_fbo);
        last_updated_time_esp = current_time;
    }


    if(ofGetFrameNum() < 1){
        depth_thread.start();
        depth_esp.start();
        depth_portrait.start();
        yolo.start();
    }

    /* write pixels to model input image */
    try{
        map_pixels.allocate(map_fbo.getWidth(), map_h, OF_PIXELS_RGB);
        map_fbo.readToPixels(map_pixels);
        model_image.setFromPixels(map_pixels);
    } catch (exception e) {
        ofLogError() << "couldn't set pixels!";
    }
    


    try{
        depth_thread.update();
        
        /* set output to fbo's */
        depth_onnx.SetPixels(model_outptut_fbo);

    } catch (exception e){
        std::cout << "Model did not run" << std::endl;
    }

    //mapper.update();
    bullet.update(server->is_active, server->getChosenNode(), nn_nodes);

    if(tsne_update_complete) {
        tsne_update_complete = false;  // Reset flag
        onTSNEUpdateComplete();  // Call your main thread function
    }

    if(ofGetElapsedTimef() > tsne_start_time + TNSE_DURATION){
        tsne_iter_idx = (tsne_iter_idx + 1) % 9;
        updateTSNEPositions(nodes);
    }
}

//--------------------------------------------------------------
void ofApp::draw(){
    ofPushStyle();
    map_fbo_alpha.begin();
    ofClear(0);
    bullet.draw();
    map_fbo_alpha.end();
    ofPopStyle();

    map_fbo.begin();
    ofClear(0, 0, 0, 0);
    map_fbo_alpha.draw(0, 0);
    map_fbo.end();

    ofTexture t = map_fbo.getTexture();
    ofTexture d = model_outptut_fbo.getTexture();

    esp_comp_fbo.begin();

    esp_shader.begin();
    esp_shader.setUniformTexture("tex0", model_esp_out_fbo.getTexture(), 0);
    model_esp_out_fbo.draw(0, 0);
    esp_comp_fbo.end();
    esp_shader.end();

    portrait_fbo.begin();

    ofBackground(0, 0, 0, 0);

    drawPortraitZ();

    portrait_fbo.end();

    /* white depth map pass */
    map_fbo_post.begin();
    ofClear(0);
    map_depth.begin();
    map_depth.setUniformTexture("tex0", d, 0);
    model_outptut_fbo.draw(0, 0);
    map_depth.end();
    map_fbo_post.end();

    drawPortrait();

    comp_fbo.begin();
    shaders.begin();
    portrait_fbo.draw(ofGetWindowWidth() / 2, 0);
    shaders.setUniformTexture("tex0", t, 0);
    shaders.setUniformTexture("tex1", bayer, 1);
    shaders.setUniform2f("resolution", map_w, map_h);
    shaders.setUniform1f("time", ofGetElapsedTimef());
    shaders.setUniform1i("frame", ofGetFrameNum());
    map_fbo_post.draw(0, 0);
    map_fbo_alpha.draw(0,0);
    shaders.end();
    comp_fbo.end();

    comp_fbo.getTexture().setTextureMinMagFilter(GL_NEAREST, GL_NEAREST);

    //mapper.draw();

    ofTranslate(0, 60);
    comp_fbo.draw(0, 0);

    if(bullet.print_debug)
        server->print();
}

void ofApp::drawPortrait(){

    portrait_pre_fbo.begin();
    ofClear(0, 0, 0, 0);

    float scale = min(
        (float)portrait_pre_fbo.getWidth() / last_chosen_node.img.getWidth(),
        (float)portrait_pre_fbo.getHeight() / last_chosen_node.img.getHeight()
    );
    float scaledWidth = last_chosen_node.img.getWidth() * scale;
    float scaledHeight = last_chosen_node.img.getHeight() * scale;
    float xPos = (portrait_pre_fbo.getWidth() - scaledWidth) / 2;
    float yPos = (portrait_pre_fbo.getHeight() - scaledHeight) / 2;

    ofPushMatrix();
    // Move to center of FBO
    ofTranslate(portrait_pre_fbo.getWidth()/2, portrait_pre_fbo.getHeight()); // 0.8
    // Move back by half the scaled image dimensions
    ofTranslate(-scaledWidth/2, -scaledHeight);
    // Draw at 0,0 since we've already translated
    last_chosen_node.img.draw(0, 0, scaledWidth, scaledHeight);
    ofPopMatrix();
    portrait_pre_fbo.end();
    
    portrait_pre_fbo_alpha.begin();
    ofPushMatrix();
    ofClear(0, 0, 0, 0);
    ofTranslate(portrait_pre_fbo.getWidth()/2, portrait_pre_fbo.getHeight());
    // Move back by half the scaled image dimensions
    ofTranslate(-scaledWidth/2, -scaledHeight);
    last_chosen_node.img.draw(0, 0, scaledWidth, scaledHeight);
    ofPopMatrix();
    portrait_pre_fbo_alpha.end();
    
    portrait_pre_fbo_alpha.readToPixels(alpha_pixels);
    alpha_demo.setFromPixels(alpha_pixels);

    float current_time = ofGetElapsedTimef();

    // we have full portrait & alpha portrait. 

    if(portrait_needs_update) {
        yolo.resetInferenceFlag();
        
        ofPixels pix;
        portrait_pre_fbo.readToPixels(pix);
        model_image_portrait.setFromPixels(pix);

        detected_faces.clear();
        yolo.update();  // This triggers the thread to run
        
        portrait_needs_update = false;
    }

    if(yolo.checkInferenceComplete()) {
        /* check if face is detected, and crop potrait */
        if(detected_faces.size() > 0){
            for(auto& f : detected_faces){
                float crop_width = f.box.x2 - f.box.x1;
                float crop_height = f.box.y2 - f.box.y1;
                
                // Calculate scaling to fit in FBO while maintaining aspect ratio
                float scale = min(
                    (float)portrait_cropped.getWidth() / crop_width,
                    (float)portrait_cropped.getHeight() / crop_height
                );
                
                float scaledWidth = crop_width * scale;
                float scaledHeight = crop_height * scale;
                
                // Calculate position to center in FBO
                float x = (portrait_cropped.getWidth() - scaledWidth) / 2;
                float y = (portrait_cropped.getHeight() - scaledHeight) / 2;
                
                portrait_cropped.begin();
                ofClear(ofColor::black);
                ofPushMatrix();
                // Move to center position
                ofTranslate(x, y);
                // Draw the cropped section at the calculated size
                model_image_portrait.drawSubsection(
                    0, 0,               // Draw at translated position
                    scaledWidth, scaledHeight,  // Draw at scaled size
                    f.box.x1, f.box.y1,        // Start crop from face detection
                    crop_width, crop_height     // Amount to crop
                );
                ofPopMatrix();
                portrait_cropped.end();
                
                portrait_cropped_alpha.begin();
                ofClear(0);
                ofPushMatrix();
                // Move to center position
                ofTranslate(x, y);
                // Draw the cropped section at the calculated size
                alpha_demo.drawSubsection(
                    0, 0,               // Draw at translated position
                    scaledWidth, scaledHeight,  // Draw at scaled size
                    f.box.x1, f.box.y1,        // Start crop from face detection
                    crop_width, crop_height     // Amount to crop
                );
                ofPopMatrix();
                portrait_cropped_alpha.end();

                ofPixels c_pix;
                portrait_cropped.readToPixels(c_pix);
                model_image_portrait_cropped.setFromPixels(c_pix);
            }
        } else {
            /* don't use cropped image for depth model */
            ofPixels c_pix;
            portrait_pre_fbo.readToPixels(c_pix);
            model_image_portrait_cropped.setFromPixels(c_pix);
        }

        /* run depth model on cropped portrait */
        depth_portrait.update();
        yolo.resetInferenceFlag();
    }

    if(depth_portrait.checkInferenceComplete()){

        depth_onnx_portrait.SetPixels(model_portrait_out_fbo);
        float planeScale = 1.0;
        int planeWidth = portrait_pre_fbo_alpha.getWidth() * planeScale;
        int planeHeight = portrait_pre_fbo_alpha.getHeight() * planeScale;
        int planeGridSize = plane_size;
        int planeColumns = planeWidth / planeGridSize;
        int planeRows = planeHeight / planeGridSize;

        plane.set(planeWidth, planeHeight, planeColumns, planeRows, OF_PRIMITIVE_TRIANGLES);
        custom_mesh = createCustomPlane(planeWidth, planeHeight, planeColumns, planeRows);
        
        /* send osc message when new node selected*/
        server->sendOSCMessage();
        depth_portrait.resetInferenceFlag();
    }
}

void ofApp::drawPortraitZ(){
    ofBackgroundGradient(ofColor::white, ofColor::orange, OF_GRADIENT_LINEAR);

    /* if no faces found use full portrait texture */
    ofTexture tex_color;

    if(detected_faces.size() > 0){
        tex_color = portrait_cropped_alpha.getTexture();
    } else {
        tex_color = portrait_pre_fbo_alpha.getTexture();
    }
    
    ofTexture tex_depth = model_portrait_out_fbo.getTexture();

    float minDepth = std::numeric_limits<float>::max();
    float maxDepth = std::numeric_limits<float>::lowest();

    ofPixels depth_pix;
    tex_depth.readToPixels(depth_pix);

    float max_r = 0;

    // Loop through only the R channel
    for(int i = 0; i < depth_pix.getWidth() * depth_pix.getHeight(); i++) {
        float depth = depth_pix[i * depth_pix.getNumChannels()] / 255.0f; // Access just R channel
        if(max_r < depth){
            max_r = depth;
        }
    }

    portrait_camera.begin();
    ofEnableDepthTest();
    
    float time = ofGetElapsedTimef() / 20;
    float p_noise_x = ofSignedNoise(time, time) * 20;
    float p_noise_y = ofSignedNoise(time + 100, time + 100) * 20;

    

    // Second plane with offset
    ofPushMatrix();
    p_depth.begin();
    p_depth.setUniform1f("gridSize", g_size);
    p_depth.setUniform2f("resolution", glm::vec2(portrait_fbo.getWidth(), portrait_fbo.getHeight()));
    p_depth.setUniform2f("threshold", glm::vec2(minDepth, max_r));
    p_depth.setUniformTexture("tex0", tex_color, 0);
    p_depth.setUniformTexture("tex1", tex_depth, 1);
    ofFill();
    // Offset the second plane - adjust these values as needed
    ofTranslate(p_noise_x - 50, p_noise_y - 0, -150);  // moved right (+200), up (+100) and back (-300)
    ofRotateDeg(60, 0, 1, 0);
    plane.draw();
    p_depth.end();
    ofPopMatrix();
    ofDisableDepthTest();

    // Second plane with offset
    ofPushMatrix();
    p_depth.begin();
    p_depth.setUniform1f("gridSize", g_size);
    p_depth.setUniform2f("resolution", glm::vec2(portrait_fbo.getWidth(), portrait_fbo.getHeight()));
    p_depth.setUniform2f("threshold", glm::vec2(minDepth, max_r));
    p_depth.setUniformTexture("tex0", tex_color, 0);
    p_depth.setUniformTexture("tex1", tex_depth, 1);
    ofFill();
    // Offset the second plane - adjust these values as needed
    ofTranslate(p_noise_x - 140, p_noise_y - 0, -150);  // moved right (+200), up (+100) and back (-300)
    ofRotateDeg(-60, 0, 1, 0);
    plane.draw();
    p_depth.end();
    ofPopMatrix();
    ofDisableDepthTest();


    ofPushMatrix();
    p_depth.begin();
    p_depth.setUniform1f("gridSize", g_size);
    p_depth.setUniform2f("resolution", glm::vec2(portrait_fbo.getWidth(), portrait_fbo.getHeight()));
    p_depth.setUniform2f("threshold", glm::vec2(minDepth, max_r));
    p_depth.setUniformTexture("tex0", tex_color, 0);
    p_depth.setUniformTexture("tex1", tex_depth, 1);
    ofFill();
    ofTranslate(p_noise_x, p_noise_y, 0);
    plane.draw();
    p_depth.end();
    ofPopMatrix();
    
    portrait_camera.end();
}

/* creates an fbo with four cropped images, to preprare for model input */
void ofApp::getNearestImages(){
    nn_nodes.clear();
    glm::vec3 cam_position = bullet.getCameraPosition();
    queryKD(cam_position, 4);

    kd_out.begin();
    ofClear(255, 255, 255, 0);

    ofImage selected_image;

    int imageIndex = 0;
    for(auto& img : esp_images){
        /* add to nn_nodes array */
        nn_nodes.push_back(&nodes[kd_result[imageIndex].second]);

        selected_image = nodes[kd_result[imageIndex].second].img;
        // Calculate the scaling factor
        float widthRatio = 128.0f / selected_image.getWidth();
        float heightRatio = 168.0f / selected_image.getHeight();
        float scale = std::max(widthRatio, heightRatio);
        
        // Calculate new dimensions
        int newWidth = std::ceil(selected_image.getWidth() * scale);
        int newHeight = std::ceil(selected_image.getHeight() * scale);
        
        // Resize the image
        selected_image.resize(newWidth, newHeight);
        
        // Calculate the crop position to center the image
        int cropX = (newWidth - 128) / 2;
        int cropY = (newHeight - 168) / 2;
        
        // Ensure crop dimensions don't exceed the image size
        int cropWidth = std::min(128, newWidth);
        int cropHeight = std::min(168, newHeight);
        
        // Calculate the drawing position based on the image index
        int drawX = (imageIndex % 2) * 128;
        int drawY = (imageIndex / 2) * 168;
        
        // Draw the resized and cropped image
        selected_image.drawSubsection(drawX, drawY, cropWidth, cropHeight, cropX, cropY);
        
        imageIndex++;
        if (imageIndex >= 4) break; // Stop after drawing 4 images
    }
    kd_out.end();

    
    kd_out.readToPixels(esp_comp_pixels);
    model_image_esp.setFromPixels(esp_comp_pixels);
}

void ofApp::createNodes(std::string json_path){

    /* read in json */
    json = ofLoadJson(json_path);

    if(!json.is_array()){
        ofLogError() << "json is not array";
        return;
    } else {
        ofLogNotice() << "json loaded";
    }

    /* setup nodes */
    bool addNode = true;  // Toggle flag to track every second node
    for(const auto& j : json) {
        if(j.contains("vector") && j["vector"].is_array() && j["error"] == 0) {
            if(addNode) {  // Only process when flag is true
                Node n;
                n.img.load(j["image"]);
                n.tex = n.img.getTexture();
                n.error = j["error"];
                std::vector<float> t_embedding;
                int count = 0;
                for (const auto& value: j["vector"]) {
                    if(value.is_number()) {
                        t_embedding.push_back(value.get<float>());
                        count++;
                    } else {
                        ofLogError() << "Vector value is not a number";
                    }
                }
                n.raw_embedding = t_embedding;
                nodes.push_back(n);
            }
            addNode = !addNode;  // Toggle the flag after each valid node
        }
    }
    ofLogNotice() << "node count: " << nodes.size();

    /* build vp_tree */
    auto vectors = createDoubleVectorFromNodes(nodes);
    vp_tree.buildTree(vectors);
    ofLogNotice() << "vpt built";

    /* run tsne */
    vector<vector<float>> tsne_input;
    std::vector<std::vector<double>> tsne_points;

    for(const auto& n : nodes){
            tsne_input.push_back(n.raw_embedding);
    }

    tsne_points = tsne.run(tsne_input, dims, perplexity, theta, normalise, runManually);

    // Option 1: Using nested loops
    for(size_t i = 0; i < 9; i++) {
        for(size_t t = 0; t <= 100; t++) {
            tsne_points = tsne.iterate();
        }
        point_iterations.push_back(tsne_points);
    }

    for(size_t i = 0; i < tsne_points.size(); i++){
        const auto& vec = point_iterations[0][i];
        auto& n = nodes[i];
        n.tsne_position = (glm::vec3(((vec[0] * 2) - 1) * tsne_scale, ((vec[1] * 2) - 1) * tsne_scale, -5.0f));
    }
}

void ofApp::updateTSNEPositions(vector<Node>& _nodes){
    if(tsne_updating) {
        ofLog() << "TSNE update already in progress";
        return;
    }
    tsne_updating = true;
    tsne_update_complete = false;  // Reset flag
    tsne_thread = std::thread(&ofApp::updateTSNEPositionsThreaded, this, std::ref(_nodes));
    tsne_thread.detach();
}

void ofApp::updateTSNEPositionsThreaded(vector<Node>& _nodes){
    ofLog() << "Starting TSNE update in thread";
    {
        std::lock_guard<std::mutex> lock(nodes_mutex);
        for(size_t i = 0; i < point_iterations[tsne_iter_idx].size(); i++) {
            const auto& vec = point_iterations[tsne_iter_idx][i];
            nodes[i].tsne_position = glm::vec3(
                ((vec[0] * 2) - 1) * tsne_scale,
                ((vec[1] * 2) - 1) * tsne_scale,
                -5.0f
            );
        }
    }
    tsne_updating = false;
    tsne_update_complete = true;  // Set completion flag
    ofLog() << "TSNE update complete";
}

std::vector<std::vector<double>> ofApp::createDoubleVectorFromNodes(const std::vector<Node>& nodes) {
    std::vector<std::vector<double>> result;
    result.reserve(nodes.size());

    for (const auto& node : nodes) {
        std::vector<double> doubleVector;
        doubleVector.reserve(node.raw_embedding.size());

        for (const float& value : node.raw_embedding) {
            doubleVector.push_back(static_cast<double>(value));
        }

        result.push_back(std::move(doubleVector));
    }

    return result;
}

void ofApp::exit(){
    std::cout << "closing server" << std::endl;
    server->close();
}

void ofApp::buildMeshes(){
    vector<ofMesh> mesh_list;
    ofLogNotice() << "building meshes: " << nodes.size();
    for(auto& n : nodes){
        mesh_list = mesh_generator.generateSimplifiedMesh(n.img);
        bullet.addMesh(mesh_list[0], mesh_list[1], n);
        mesh_list.clear();
    }
    ofLogNotice() << "finished building meshes";
}

void ofApp::buildKDTree(){
    pointVec kd_points;

    for(const auto& n : nodes){

        glm::vec3 n_pos = n.collider->getPosition();

        std::vector< double > p = {
            static_cast<double>(n_pos.x),
            static_cast<double>(n_pos.y)
        };

        kd_points.push_back(p);
    }

    kd_tree = std::make_unique<KDTree>(kd_points);

    kd_result.resize(4);
}

void ofApp::queryKD(glm::vec3& _position, int k){

    kd_result.clear();    

    vector<double> kd_input = { 
        static_cast<double>(_position.x),
        static_cast<double>(_position.y)
    };

    auto res = kd_tree->nearest_pointIndices(kd_input, k);

    for (const auto& r : res) {
        kd_result.push_back(r);
    }

    size_t index = kd_result[0].second;

    //std::cout << "camera positions: " << _position << ", nearest node: "<< nodes[index].tsne_position << " " << index << std::endl;

}

void ofApp::keyPressed(int key) {
    switch(key) {
        case OF_KEY_UP:
            g_size += 5;
            ofLog() << "Grid Size: " << g_size;
            break;
        case OF_KEY_DOWN:
            g_size -= 5;
            ofLog() << "Grid Size: " << g_size;
            break;
        case 'q':
            plane_size += 5;
            ofLog() << "Plane Size: " << plane_size;
            break;
        case 'w':
            plane_size -=5;
            ofLog() << "Plane Size: " << plane_size;
            break;
        case 'c':
            // cam_pos_idx++;
            // if(cam_pos_idx > cam_positions.size() - 1)
            //     cam_pos_idx = 0;
            portrait_camera.setPosition(cam_positions[cam_pos_idx].position);
            portrait_camera.setOrientation(cam_positions[cam_pos_idx].rotation);
            portrait_camera.setScale(cam_positions[cam_pos_idx].scale);
        case 'p':
            if(bullet.print_debug){
                bullet.print_debug = false;
            } else {
                bullet.print_debug = true;
            }
            std::cout << portrait_camera.getPosition() << " : " <<portrait_camera.getOrientationQuat()<< " : " << portrait_camera.getScale() << std::endl;
            break;
        case 't':
            tsne_iter_idx++;
            updateTSNEPositions(nodes);
    }
    //mapper.keyPressed(key);
}

// Your main thread function:
void ofApp::onTSNEUpdateComplete() {
    // Do whatever you need to do after TSNE update
    // This will run on the main thread
    bullet.updateTSNEPosition(nodes);
    tsne_start_time = ofGetElapsedTimef();
    ofLog() << "Handling TSNE update completion on main thread";
}

void ofApp::keyReleased(int key){
    //mapper.keyReleased(key);
}

//--------------------------------------------------------------
void ofApp::mouseDragged(int x, int y, int button){
    //mapper.mouseDragged(x, y, button);
}

//--------------------------------------------------------------
void ofApp::mousePressed(int x, int y, int button){
    //mapper.mousePressed(x, y, button);
}

//--------------------------------------------------------------
void ofApp::mouseReleased(int x, int y, int button){
    //mapper.mouseReleased(x, y, button);
}

ofMesh ofApp::createCustomPlane(float width, float height, int numX, int numY) {
    ofMesh mesh;
    mesh.setMode(OF_PRIMITIVE_TRIANGLES);
    
    // Create vertices with non-uniform distribution
    for(int y = 0; y <= numY; y++) {
        for(int x = 0; x <= numX; x++) {
            // Get base uniform position
            float xPos = (float)x/numX * width;
            float yPos = (float)y/numY * height;

            float texX = (float)x/numX;
            float texY = 1.0 - (float)y/numY;
            
            // Add random offset but keep edges fixed
            if(x != 0 && x != numX && y != 0 && y != numY) {
                // Random offset proportional to cell size
                float cellWidth = width/numX;
                float cellHeight = height/numY;
                
                xPos += ofRandom(-cellWidth * 0.2, cellWidth * 0.2);
                yPos += ofRandom(-cellHeight * 0.2, cellHeight * 0.2);
            }
            
            mesh.addVertex(glm::vec3(xPos, yPos, 0));
            
            // Add texture coordinates (normalized)
            mesh.addTexCoord(glm::vec2(texX, texY));
            mesh.addColor(ofColor(255, 255, 255, 255));
        }
    }
    
    // Create triangles
    for(int y = 0; y < numY; y++) {
        for(int x = 0; x < numX; x++) {
            // Get indices for corners of quad
            int nw = y * (numX + 1) + x;
            int ne = nw + 1;
            int sw = nw + (numX + 1);
            int se = sw + 1;
            
            // Add triangles
            mesh.addIndex(nw);
            mesh.addIndex(ne);
            mesh.addIndex(se);
            
            mesh.addIndex(nw);
            mesh.addIndex(se);
            mesh.addIndex(sw);
        }
    }
    
    return mesh;
}