// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include <iostream>
#include <fstream>
#include <string>
#include <iomanip>

#include <k4a/k4a.h>
#include <k4arecord/playback.h>
#include <k4abt.h>
#include <nlohmann/json.hpp>

#include <BodyTrackingHelpers.h>
#include <Utilities.h>

using namespace std;
using namespace nlohmann;

bool predict_joints(json &frames_json, int frame_count, k4abt_tracker_t tracker, k4a_capture_t capture_handle)
{
    k4a_wait_result_t queue_capture_result = k4abt_tracker_enqueue_capture(tracker, capture_handle, K4A_WAIT_INFINITE);
    if (queue_capture_result != K4A_WAIT_RESULT_SUCCEEDED)
    {
        cerr << "Error! Adding capture to tracker process queue failed!" << endl;
        return false;
    }

    k4abt_frame_t body_frame = nullptr;
    k4a_wait_result_t pop_frame_result = k4abt_tracker_pop_result(tracker, &body_frame, K4A_WAIT_INFINITE);
    if (pop_frame_result != K4A_WAIT_RESULT_SUCCEEDED)
    {
        cerr << "Error! Popping body tracking result failed!" << endl;
        return false;
    }

    uint32_t num_bodies = k4abt_frame_get_num_bodies(body_frame);
    uint64_t timestamp = k4abt_frame_get_device_timestamp_usec(body_frame);

    json frame_result_json;
    frame_result_json["timestamp_usec"] = timestamp;
    frame_result_json["frame_id"] = frame_count;
    frame_result_json["num_bodies"] = num_bodies;
    frame_result_json["bodies"] = json::array();
    for (uint32_t i = 0; i < num_bodies; i++)
    {
        k4abt_skeleton_t skeleton;
        VERIFY(k4abt_frame_get_body_skeleton(body_frame, i, &skeleton), "Get body from body frame failed!");
        json body_result_json;
        int body_id = k4abt_frame_get_body_id(body_frame, i);
        body_result_json["body_id"] = body_id;

        for (int j = 0; j < (int)K4ABT_JOINT_COUNT; j++)
        {
            body_result_json["joint_positions"].push_back( {    skeleton.joints[j].position.xyz.x,
                                                                skeleton.joints[j].position.xyz.y,
                                                                skeleton.joints[j].position.xyz.z });

            body_result_json["joint_orientations"].push_back({  skeleton.joints[j].orientation.wxyz.w,
                                                                skeleton.joints[j].orientation.wxyz.x,
                                                                skeleton.joints[j].orientation.wxyz.y,
                                                                skeleton.joints[j].orientation.wxyz.z });
        }
        frame_result_json["bodies"].push_back(body_result_json);
    }
    frames_json.push_back(frame_result_json);
    k4abt_frame_release(body_frame);

    return true;
}

bool check_depth_image_exists(k4a_capture_t capture)
{
    k4a_image_t depth = k4a_capture_get_depth_image(capture);
    if (depth != nullptr)
    {
        k4a_image_release(depth);
        return true;
    }
    else
    {
        return false;
    }
}

bool process_mkv_offline(const char* input_path, const char* output_path, k4abt_tracker_configuration_t tracker_config = K4ABT_TRACKER_CONFIG_DEFAULT )
{
    k4a_playback_t playback_handle = nullptr;
    k4a_result_t result = k4a_playback_open(input_path, &playback_handle);
    if (result != K4A_RESULT_SUCCEEDED)
    {
        cerr << "Cannot open recording at " << input_path << endl;
        return false;
    }

    k4a_calibration_t calibration;
    result = k4a_playback_get_calibration(playback_handle, &calibration);
    if (result != K4A_RESULT_SUCCEEDED)
    {
        cerr << "Failed to get calibration" << endl;
        return false;
    }

    k4abt_tracker_t tracker = NULL;
    if (K4A_RESULT_SUCCEEDED != k4abt_tracker_create(&calibration, tracker_config, &tracker))
    {
        cerr << "Body tracker initialization failed!" << endl;
        return false;
    }

    json json_output;
    json_output["k4abt_sdk_version"] = K4ABT_VERSION_STR;
    json_output["source_file"] = input_path;

    // Store all joint names to the json
    json_output["joint_names"] = json::array();
    for (int i = 0; i < (int)K4ABT_JOINT_COUNT; i++)
    {
        json_output["joint_names"].push_back(g_jointNames.find((k4abt_joint_id_t)i)->second);
    }

    // Store all bone linkings to the json
    json_output["bone_list"] = json::array();
    for (int i = 0; i < (int)g_boneList.size(); i++)
    {
        json_output["bone_list"].push_back({ g_jointNames.find(g_boneList[i].first)->second,
                                             g_jointNames.find(g_boneList[i].second)->second });
    }

    cout << "Tracking " << input_path << endl;

    int frame_count = 0;
    json frames_json = json::array();
    bool success = true;
    while (true)
    {
        k4a_capture_t capture_handle = nullptr;
        k4a_stream_result_t stream_result = k4a_playback_get_next_capture(playback_handle, &capture_handle);
        if (stream_result == K4A_STREAM_RESULT_EOF)
        {
            break;
        }

        cout << "frame " << frame_count << '\r';
        if (stream_result == K4A_STREAM_RESULT_SUCCEEDED)
        {
            // Only try to predict joints when capture contains depth image
            if (check_depth_image_exists(capture_handle))
            {
                success = predict_joints(frames_json, frame_count, tracker, capture_handle);
                k4a_capture_release(capture_handle);
                if (!success)
                {
                    cerr << "Predict joints failed for clip at frame " << frame_count << endl;
                    break;
                }
            }
        }
        else
        {
            success = false;
            cerr << "Stream error for clip at frame " << frame_count << endl;
            break;
        }

        frame_count++;
    }

    if (success)
    {
        json_output["frames"] = frames_json;
        cout << endl << "DONE " << endl;

        cout << "Total read " << frame_count << " frames" << endl;
        std::ofstream output_file(output_path);
        output_file << std::setw(4) << json_output << std::endl;
        cout << "Results saved in " << output_path;
    }

    k4abt_tracker_shutdown(tracker);
    k4abt_tracker_destroy(tracker);
    k4a_playback_close(playback_handle);

    return success;
}

void PrintUsage()
{
#ifdef _WIN32
    cout << "Usage: k4abt_offline_processor.exe <input_mkv_file> <output_json_file> [processing_mode] [-model MODEL_FILE_PATH]\n\t[Optional] processing_mode\n\t\tCPU\n\t\tCUDA\n\t\tTensorRT\n\t\tDirectML ( default )" << endl;
#else
    cout << "Usage: k4abt_offline_processor.exe <input_mkv_file> <output_json_file> [processing_mode] [-model MODEL_FILE_PATH]\n\t[Optional] processing_mode\n\t\tCPU\n\t\tCUDA ( default )\n\t\tTensorRT" << endl;
#endif
}

bool ProcessArguments(k4abt_tracker_configuration_t &tracker_config, int argc, char** argv)
{
    if (argc < 3)
    {
        PrintUsage();
        return false;
    }
    for( int i = 3; i < argc; i++ )
    {
        if (0 == strcmp(argv[i], "TensorRT"))
        {
            tracker_config.processing_mode = K4ABT_TRACKER_PROCESSING_MODE_GPU_TENSORRT;
        }
        else if (0 == strcmp(argv[i], "CUDA"))
        {
            tracker_config.processing_mode = K4ABT_TRACKER_PROCESSING_MODE_GPU_CUDA;
        }
        else if (0 == strcmp(argv[i], "CPU"))
        {
            tracker_config.processing_mode = K4ABT_TRACKER_PROCESSING_MODE_CPU;
        }
#ifdef _WIN32
        else if (0 == strcmp(argv[i], "DirectML"))
        {
            tracker_config.processing_mode = K4ABT_TRACKER_PROCESSING_MODE_GPU_DIRECTML;
        }
#endif
        else if (0 == strcmp(argv[i], "-model"))
        {
            if (i < argc - 1)
                tracker_config.model_path = argv[++i];
            else
            {
                printf("Error: model filepath missing\n");
                PrintUsage();
                return false;
            }
        }
        else
        {
            PrintUsage();
            return false;
        }
    }
    return true;
}

int main(int argc, char **argv)
{
    k4abt_tracker_configuration_t tracker_config = K4ABT_TRACKER_CONFIG_DEFAULT;
    if (!ProcessArguments(tracker_config, argc, argv))
        return -1;
    return process_mkv_offline(argv[1], argv[2], tracker_config) ? 0 : -1;
}