import json
from pathlib import Path

import numpy as np
import pytest

import rpt

ROOT = Path(__file__).resolve().parents[1]

JOINT_NAMES = [
    "nose",
    "eye_left",
    "eye_right",
    "ear_left",
    "ear_right",
    "shoulder_left",
    "shoulder_right",
    "elbow_left",
    "elbow_right",
    "wrist_left",
    "wrist_right",
    "hip_left",
    "hip_right",
    "knee_left",
    "knee_right",
    "ankle_left",
    "ankle_right",
    "hip_middle",
    "shoulder_middle",
    "head",
]


def load_case(camera_path: str, pose_path: str):
    with (ROOT / camera_path).open("r", encoding="utf-8") as file:
        camera_data = json.load(file)
    with (ROOT / pose_path).open("r", encoding="utf-8") as file:
        pose_data = json.load(file)

    poses_2d, person_counts = rpt.pack_poses_2d(pose_data["2D"], joint_count=len(JOINT_NAMES))
    cameras = rpt.convert_cameras(camera_data["cameras"])
    return poses_2d, person_counts, cameras


def test_camera_structure_repr():
    camera = rpt.Camera()
    camera.name = "Camera 1"
    camera.K = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    camera.DC = [0, 0, 0, 0, 0]
    camera.R = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    camera.T = [[1], [2], [3]]
    camera.width = 640
    camera.height = 480
    camera.type = "pinhole"

    rendered = repr(camera)
    assert "Camera 1" in rendered
    assert "pinhole" in rendered


@pytest.mark.parametrize(
    ("camera_path", "pose_path", "roomparams"),
    [
        ("data/h1/sample.json", "tests/poses_h1.json", [[4.8, 6.0, 2.0], [0.0, 0.0, 1.0]]),
        ("data/p1/sample.json", "tests/poses_p1.json", [[5.6, 6.4, 2.4], [0.0, -0.5, 1.2]]),
        ("data/e1/sample.json", "tests/poses_e1.json", [[6.0, 5.0, 2.0], [1.5, 1.0, -0.5]]),
    ],
)
def test_triangulate_samples(camera_path: str, pose_path: str, roomparams):
    poses_2d, person_counts, cameras = load_case(camera_path, pose_path)
    poses_3d = rpt.triangulate_poses(
        poses_2d,
        person_counts,
        cameras,
        np.asarray(roomparams, dtype=np.float32),
        JOINT_NAMES,
        min_match_score=0.95,
    )

    assert isinstance(poses_3d, np.ndarray)
    assert poses_3d.dtype == np.float32
    assert poses_3d.ndim == 3
    assert poses_3d.shape[1:] == (len(JOINT_NAMES), 4)
    assert poses_3d.shape[0] > 0
    assert np.isfinite(poses_3d).all()


def test_triangulate_repeatability():
    poses_2d, person_counts, cameras = load_case("data/p1/sample.json", "tests/poses_p1.json")
    roomparams = np.asarray([[5.6, 6.4, 2.4], [0.0, -0.5, 1.2]], dtype=np.float32)

    first = rpt.triangulate_poses(
        poses_2d, person_counts, cameras, roomparams, JOINT_NAMES, min_match_score=0.95
    )
    second = rpt.triangulate_poses(
        poses_2d, person_counts, cameras, roomparams, JOINT_NAMES, min_match_score=0.95
    )

    np.testing.assert_allclose(first, second, rtol=1e-5, atol=1e-5)


def test_triangulate_does_not_mutate_inputs():
    poses_2d, person_counts, cameras = load_case("data/h1/sample.json", "tests/poses_h1.json")
    roomparams = np.asarray([[4.8, 6.0, 2.0], [0.0, 0.0, 1.0]], dtype=np.float32)

    poses_before = poses_2d.copy()
    counts_before = person_counts.copy()

    rpt.triangulate_poses(poses_2d, person_counts, cameras, roomparams, JOINT_NAMES)

    np.testing.assert_array_equal(poses_2d, poses_before)
    np.testing.assert_array_equal(person_counts, counts_before)


def test_pack_poses_2d_from_ragged_inputs():
    packed, counts = rpt.pack_poses_2d(
        [
            [[[1, 2, 0.5], [3, 4, 0.6]]],
            np.asarray(
                [
                    [[5, 6, 0.7], [7, 8, 0.8]],
                    [[9, 10, 0.9], [11, 12, 1.0]],
                ],
                dtype=np.float32,
            ),
            [],
        ],
        joint_count=2,
    )

    assert packed.shape == (3, 2, 2, 3)
    assert packed.dtype == np.float32
    np.testing.assert_array_equal(counts, np.asarray([1, 2, 0], dtype=np.uint32))
    np.testing.assert_allclose(packed[0, 0], np.asarray([[1, 2, 0.5], [3, 4, 0.6]], dtype=np.float32))
    np.testing.assert_allclose(packed[1, 1], np.asarray([[9, 10, 0.9], [11, 12, 1.0]], dtype=np.float32))
    np.testing.assert_array_equal(packed[2], np.zeros((2, 2, 3), dtype=np.float32))


def test_pack_poses_2d_rejects_inconsistent_joint_count():
    with pytest.raises(ValueError, match="joint count"):
        rpt.pack_poses_2d(
            [
                [[[1, 2, 0.5], [3, 4, 0.6]]],
                [[[5, 6, 0.7]]],
            ]
        )


def test_pack_poses_2d_rejects_invalid_last_dimension():
    with pytest.raises(ValueError, match="shape"):
        rpt.pack_poses_2d([np.zeros((1, 2, 2), dtype=np.float32)])