#pragma once

#include <cstdint>
#include <memory>
#include <string>
#include <vector>

#include "camera.hpp"

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

using RaggedPoses2D = std::vector<std::vector<std::vector<std::array<float, 3>>>>;
using NestedPoses3D = std::vector<std::vector<std::array<float, 4>>>;

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

struct PoseBatch2D
{
    std::vector<float> data;
    std::vector<uint32_t> person_counts;
    size_t num_views = 0;
    size_t max_persons = 0;
    size_t num_joints = 0;

    float &at(size_t view, size_t person, size_t joint, size_t coord);
    const float &at(size_t view, size_t person, size_t joint, size_t coord) const;

    RaggedPoses2D to_nested() const;
    static PoseBatch2D from_nested(const RaggedPoses2D &poses_2d);
};

struct PoseBatch3D
{
    std::vector<float> data;
    size_t num_persons = 0;
    size_t num_joints = 0;

    float &at(size_t person, size_t joint, size_t coord);
    const float &at(size_t person, size_t joint, size_t coord) const;

    NestedPoses3D to_nested() const;
    static PoseBatch3D from_nested(const NestedPoses3D &poses_3d);
};

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

class TriangulatorInternal;

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

class Triangulator
{
public:
    /**
     * Triangulator to predict poses from multiple views.
     *
     * @param min_match_score Minimum score to consider a triangulated joint as valid.
     * @param min_group_size Minimum number of camera pairs that need to see a person.
     */
    Triangulator(
        float min_match_score = 0.95,
        size_t min_group_size = 1);

    ~Triangulator();

    /**
     * Calculate a triangulation using a padded pose tensor.
     *
     * @param poses_2d Padded poses of shape [views, max_persons, joints, 3].
     * @param cameras List of cameras.
     * @param roomparams Room parameters (room size, room center).
     * @param joint_names List of 2D joint names.
     *
     * @return Pose tensor of shape [persons, joints, 4].
     */
    PoseBatch3D triangulate_poses(
        const PoseBatch2D &poses_2d,
        const std::vector<Camera> &cameras,
        const std::array<std::array<float, 3>, 2> &roomparams,
        const std::vector<std::string> &joint_names);

    /**
     * Compatibility overload for callers that still produce ragged nested vectors.
     */
    NestedPoses3D triangulate_poses(
        const RaggedPoses2D &poses_2d,
        const std::vector<Camera> &cameras,
        const std::array<std::array<float, 3>, 2> &roomparams,
        const std::vector<std::string> &joint_names);

    /** Reset the triangulator. */
    void reset();

    /** Print triangulation statistics. */
    void print_stats();

private:
    std::unique_ptr<TriangulatorInternal> triangulator;
};