#pragma once

#include <array>
#include <string>
#include <vector>

#include "camera.hpp"

// =================================================================================================

class CameraInternal
{
public:
    CameraInternal(const Camera &cam);

    Camera cam;

    std::array<std::array<float, 3>, 3> invK;
    std::array<std::array<float, 3>, 3> invR;
    std::array<float, 3> center;

    static std::array<std::array<float, 3>, 3> transpose3x3(
        const std::array<std::array<float, 3>, 3> &M);

    static std::array<std::array<float, 3>, 3> invert3x3(
        const std::array<std::array<float, 3>, 3> &M);

    static void undistort_point_pinhole(std::array<float, 3> &p, const std::vector<float> &k);
    static void undistort_point_fisheye(std::array<float, 3> &p, const std::vector<float> &k);

    std::array<std::array<float, 3>, 3> calc_optimal_camera_matrix_fisheye(
        float balance, std::pair<int, int> new_size);
    std::array<std::array<float, 3>, 3> calc_optimal_camera_matrix_pinhole(
        float alpha, std::pair<int, int> new_size);
};

// =================================================================================================

class TriangulatorInternal
{
public:
    TriangulatorInternal(float min_match_score, size_t min_group_size);

    std::vector<std::vector<std::array<float, 4>>> triangulate_poses(
        const std::vector<std::vector<std::vector<std::array<float, 3>>>> &poses_2d,
        const std::vector<Camera> &cameras,
        const std::array<std::array<float, 3>, 2> &roomparams,
        const std::vector<std::string> &joint_names);

    void reset();
    void print_stats();

private:
    float min_match_score;
    float min_group_size;

    const std::vector<std::string> core_joints = {
        "shoulder_left",
        "shoulder_right",
        "hip_left",
        "hip_right",
        "elbow_left",
        "elbow_right",
        "knee_left",
        "knee_right",
        "wrist_left",
        "wrist_right",
        "ankle_left",
        "ankle_right",
    };
    const std::vector<std::pair<std::string, std::string>> core_limbs = {
        {"knee_left", "ankle_left"},
        {"hip_left", "knee_left"},
        {"hip_right", "knee_right"},
        {"knee_right", "ankle_right"},
        {"elbow_left", "wrist_left"},
        {"elbow_right", "wrist_right"},
        {"shoulder_left", "elbow_left"},
        {"shoulder_right", "elbow_right"},
    };
    std::vector<std::vector<std::array<float, 4>>> last_poses_3d;

    void undistort_poses(
        std::vector<std::vector<std::array<float, 3>>> &poses_2d, CameraInternal &icam);

    std::tuple<std::vector<std::vector<std::array<float, 3>>>, std::vector<std::vector<float>>>
    project_poses(
        const std::vector<std::vector<std::array<float, 4>>> &poses_3d,
        const CameraInternal &icam,
        bool calc_dists);

    float calc_pose_score(
        const std::vector<std::array<float, 3>> &pose1,
        const std::vector<std::array<float, 3>> &pose2,
        const std::vector<float> &dist1,
        const CameraInternal &icam);

    std::vector<float> score_projection(
        const std::vector<std::array<float, 3>> &pose,
        const std::vector<std::array<float, 3>> &repro,
        const std::vector<float> &dists,
        const std::vector<bool> &mask,
        float iscale);

    std::pair<std::vector<std::array<float, 4>>, float> triangulate_and_score(
        const std::vector<std::array<float, 3>> &pose1,
        const std::vector<std::array<float, 3>> &pose2,
        const CameraInternal &cam1,
        const CameraInternal &cam2,
        const std::array<std::array<float, 3>, 2> &roomparams,
        const std::vector<std::array<size_t, 2>> &core_limbs_idx);

    std::vector<std::tuple<
        std::array<float, 3>, std::vector<std::array<float, 4>>, std::vector<int>>>
    calc_grouping(
        const std::vector<std::pair<
            std::tuple<int, int, int, int>, std::pair<int, int>>> &all_pairs,
        const std::vector<std::pair<std::vector<std::array<float, 4>>, float>> &all_scored_poses,
        float min_score);

    std::vector<std::array<float, 4>> merge_group(
        const std::vector<std::vector<std::array<float, 4>>> &poses_3d,
        float min_score);

    void add_extra_joints(
        std::vector<std::vector<std::array<float, 4>>> &poses,
        const std::vector<std::string> &joint_names);

    void filter_poses(
        std::vector<std::vector<std::array<float, 4>>> &poses,
        const std::array<std::array<float, 3>, 2> &roomparams,
        const std::vector<size_t> &core_joint_idx,
        const std::vector<std::array<size_t, 2>> &core_limbs_idx);

    void add_missing_joints(
        std::vector<std::vector<std::array<float, 4>>> &poses,
        const std::vector<std::string> &joint_names);

    // Statistics
    double num_calls = 0;
    double total_time = 0;
    double init_time = 0;
    double undistort_time = 0;
    double project_time = 0;
    double match_time = 0;
    double pairs_time = 0;
    double pair_scoring_time = 0;
    double grouping_time = 0;
    double full_time = 0;
    double merge_time = 0;
    double post_time = 0;
    double convert_time = 0;
};