"""
Process images test
"""
import os
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
from nerfstudio.data.utils.colmap_parsing_utils import (
    Camera,
    Image as ColmapImage,
    Point3D,
    qvec2rotmat,
    write_cameras_binary,
    write_images_binary,
    write_points3D_binary,
)
from nerfstudio.process_data.images_to_nerfstudio_dataset import ImagesToNerfstudioDataset


def random_quaternion(num_poses: int):
    """
    Generates random rotation quaternion.
    """
    u, v, w = np.random.uniform(size=(3, num_poses))
    quaternion = np.stack(
        (
            np.sqrt(1 - u) * np.sin(2 * np.pi * v),
            np.sqrt(1 - u) * np.cos(2 * np.pi * v),
            np.sqrt(u) * np.sin(2 * np.pi * w),
            np.sqrt(u) * np.cos(2 * np.pi * w),
        ),
        -1,
    )
    return quaternion


def test_process_images_skip_colmap(tmp_path: Path):
    """
    Test ns-process-data images
    """
    # Mock a colmap sparse model
    width = 100
    height = 150
    sparse_path = tmp_path / "sparse" / "0"
    sparse_path.mkdir(exist_ok=True, parents=True)
    (tmp_path / "images").mkdir(exist_ok=True, parents=True)
    write_cameras_binary(
        {1: Camera(1, "OPENCV", width, height, [110, 110, 50, 75, 0, 0, 0, 0, 0, 0])},
        sparse_path / "cameras.bin",
    )
    write_points3D_binary(
        {
            1: Point3D(
                id=1,
                xyz=np.array([0, 0, 0]),
                rgb=np.array([0, 0, 0]),
                error=np.array([0]),
                image_ids=np.array([1]),
                point2D_idxs=np.array([0]),
            ),
        },
        sparse_path / "points3D.bin",
    )
    frames = {}
    num_frames = 10
    qvecs = random_quaternion(num_frames)
    tvecs = np.random.uniform(size=(num_frames, 3))
    original_poses = np.concatenate(
        (
            np.concatenate(
                (
                    np.stack(list(map(qvec2rotmat, qvecs))),
                    tvecs[:, :, None],
                ),
                -1,
            ),
            np.array([[[0, 0, 0, 1]]], dtype=qvecs.dtype).repeat(num_frames, 0),
        ),
        -2,
    )
    for i in range(num_frames):
        frames[i + 1] = ColmapImage(i + 1, qvecs[i], tvecs[i], 1, f"image_{i}.png", [], [])
        Image.new("RGB", (width, height)).save(tmp_path / "images" / f"image_{i}.png")
    write_images_binary(frames, sparse_path / "images.bin")

    # Mock missing COLMAP and ffmpeg in the dev env
    old_path = os.environ.get("PATH", "")
    os.environ["PATH"] = str(tmp_path / "mocked_bin") + f":{old_path}"
    (tmp_path / "mocked_bin").mkdir()
    (tmp_path / "mocked_bin" / "colmap").touch(mode=0o777)
    (tmp_path / "mocked_bin" / "ffmpeg").touch(mode=0o777)

    # Convert images into a NerfStudio dataset
    cmd = ImagesToNerfstudioDataset(
        data=tmp_path / "images", output_dir=tmp_path / "nerfstudio", colmap_model_path=sparse_path, skip_colmap=True
    )
    cmd.main()
    os.environ["PATH"] = old_path

    assert (tmp_path / "nerfstudio" / "transforms.json").exists()
    parser = NerfstudioDataParserConfig(
        data=tmp_path / "nerfstudio",
        downscale_factor=None,
        orientation_method="none",  # orientation_method,
        center_method="none",
        auto_scale_poses=False,
    ).setup()
    outputs = parser.get_dataparser_outputs("train")
    assert len(outputs.image_filenames) == 9
    assert torch.is_tensor(outputs.dataparser_transform)

    # Test if the original poses can be obtained back
    gt_poses = original_poses[[0, 1, 2, 3, 4, 5, 6, 7, 9]]
    dataparser_poses = outputs.transform_poses_to_original_space(outputs.cameras.camera_to_worlds, "opencv").numpy()
    dataparser_poses = np.concatenate(
        (dataparser_poses, np.array([[[0, 0, 0, 1]]]).repeat(len(dataparser_poses), 0)), 1
    )
    dataparser_poses = np.linalg.inv(dataparser_poses)
    np.testing.assert_allclose(gt_poses, dataparser_poses, rtol=0, atol=1e-5)