#include <algorithm>
#include <cstdint>
#include <stdexcept>

#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
#include <nanobind/stl/array.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/vector.h>

#include "interface.hpp"

namespace nb = nanobind;
using namespace nb::literals;

namespace
{
using PoseArray2D =
    nb::ndarray<nb::numpy, const float, nb::shape<-1, -1, -1, 3>, nb::c_contig>;
using CountArray = nb::ndarray<nb::numpy, const uint32_t, nb::shape<-1>, nb::c_contig>;
using RoomArray = nb::ndarray<nb::numpy, const float, nb::shape<2, 3>, nb::c_contig>;
using PoseArray3D = nb::ndarray<nb::numpy, float, nb::shape<-1, -1, 4>, nb::c_contig>;

PoseBatch2DView pose_batch_view_from_numpy(const PoseArray2D &poses_2d, const CountArray &person_counts)
{
    if (poses_2d.shape(0) != person_counts.shape(0))
    {
        throw std::invalid_argument("poses_2d and person_counts must have the same number of views.");
    }

    for (size_t i = 0; i < static_cast<size_t>(person_counts.shape(0)); ++i)
    {
        if (person_counts(i) > poses_2d.shape(1))
        {
            throw std::invalid_argument("person_counts entries must not exceed the padded person dimension.");
        }
    }

    return PoseBatch2DView {
        poses_2d.data(),
        person_counts.data(),
        static_cast<size_t>(poses_2d.shape(0)),
        static_cast<size_t>(poses_2d.shape(1)),
        static_cast<size_t>(poses_2d.shape(2)),
    };
}

std::array<std::array<float, 3>, 2> roomparams_from_numpy(const RoomArray &roomparams)
{
    std::array<std::array<float, 3>, 2> result {};
    for (size_t i = 0; i < 2; ++i)
    {
        for (size_t j = 0; j < 3; ++j)
        {
            result[i][j] = roomparams(i, j);
        }
    }
    return result;
}

PoseArray3D pose_batch_to_numpy(PoseBatch3D batch)
{
    auto *storage = new std::vector<float>(std::move(batch.data));
    nb::capsule owner(storage, [](void *value) noexcept
    {
        delete static_cast<std::vector<float> *>(value);
    });

    const size_t shape[3] = {batch.num_persons, batch.num_joints, 4};
    return PoseArray3D(storage->data(), 3, shape, owner);
}
} // namespace

NB_MODULE(_core, m)
{
    nb::class_<Camera>(m, "Camera")
        .def(nb::init<>())
        .def_rw("name", &Camera::name)
        .def_rw("K", &Camera::K)
        .def_rw("DC", &Camera::DC)
        .def_rw("R", &Camera::R)
        .def_rw("T", &Camera::T)
        .def_rw("width", &Camera::width)
        .def_rw("height", &Camera::height)
        .def_rw("type", &Camera::type)
        .def("__repr__", [](const Camera &camera)
        {
            return camera.to_string();
        });

    m.def(
        "triangulate_poses",
        [](const PoseArray2D &poses_2d,
           const CountArray &person_counts,
           const std::vector<Camera> &cameras,
           const RoomArray &roomparams,
           const std::vector<std::string> &joint_names,
           float min_match_score,
           size_t min_group_size)
        {
            const PoseBatch2DView pose_batch = pose_batch_view_from_numpy(poses_2d, person_counts);
            const auto room = roomparams_from_numpy(roomparams);
            PoseBatch3D poses_3d = triangulate_poses(
                pose_batch, cameras, room, joint_names, min_match_score, min_group_size);
            return pose_batch_to_numpy(std::move(poses_3d));
        },
        "poses_2d"_a,
        "person_counts"_a,
        "cameras"_a,
        "roomparams"_a,
        "joint_names"_a,
        "min_match_score"_a = 0.95f,
        "min_group_size"_a = 1);
}