#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);
    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);

    map_pixels.allocate(map_fbo.getWidth(), map_h, OF_PIXELS_RGB);

    /* 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";

    /* bullet setup */
    bullet.setup(nodes);

    /* 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_Small, 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, &model_portrait_out_fbo, &depth_onnx_portrait);

    /* camera settings for portrait */
    CameraPosition cp;
    cp.position = glm::vec3(116.388, -441.054, 967.724);
    cp.rotation = glm::quat(0.9155, 0.3170, 0.02, -0.035);
    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(0.5);
    // 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();
}

//--------------------------------------------------------------
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 >= 3){
        buildKDTree();
        getNearestImages();
        depth_esp.update();
        depth_onnx_esp.SetPixels(model_esp_out_fbo);
        last_updated_time = current_time;
    }
    


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

    /* write pixels to model input image */
    try{
        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());
    
    //std::cout << portrait_camera.getPosition() << " : " <<portrait_camera.getOrientationQuat()<< std::endl;
}

//--------------------------------------------------------------
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();
    ofClear(ofColor::lightGrey); //uncommet when changing to map
    ofBackgroundGradient(ofColor::darkGrey, ofColor::whiteSmoke, OF_GRADIENT_LINEAR);
    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);
    /* actual map */
    map_fbo_alpha.draw(0,0);
    //model_portrait_out_fbo.draw(0,0);
    //portrait_pre_fbo_alpha.draw(model_portrait_out_fbo.getWidth(),0);
    shaders.end();
    comp_fbo.end();

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

    //mapper.draw();
    comp_fbo.draw(0,60);

    //server->print();
}

void ofApp::drawPortrait(){

    float p_scale =  1 + ( ((1 + ofNoise(ofGetElapsedTimef() / 5)) / 2) * 0.5);

    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
    // Apply scale
    ofScale(1);
    // 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());
    // Apply scale
    ofScale(1);
    // 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();

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

        // Only run portrait model when needed
        depth_portrait.update();
        depth_onnx_portrait.SetPixels(model_portrait_out_fbo);

        /* create mesh */
        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);
        plane.mapTexCoords(0, 0, planeWidth, planeHeight);
        custom_mesh = createCustomPlane(planeWidth, planeHeight, planeColumns, planeRows);
       // ofLog() << "new mesh";
    }
}

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

    ofTexture 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() / 10;
    float p_noise_x = ofSignedNoise(time, time) * 20;
    float p_noise_y = ofSignedNoise(time + 100, time + 100) * 20;

    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);
    custom_mesh.draw();
    p_depth.end();
    ofPopMatrix();
    ofDisableDepthTest();
    
    portrait_camera.end();
}

/* creates an fbo with four cropped images, to preprare for model input */
void ofApp::getNearestImages(){
    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){
        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);

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

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);
    }
    //mapper.keyPressed(key);
}

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.4, cellWidth * 0.4);
                yPos += ofRandom(-cellHeight * 0.4, cellHeight * 0.4);
            }
            
            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;
}