#!/usr/bin/env python3

import numpy as np

from bornagain.ba_view3d import (
    _arrow_transform,
    _beam_endpoints,
    _contrasting_text_color,
    _viewer_kwargs,
)


def test_beams_are_enabled_by_default():
    """Tests subprocess-safe propagation of the beam display option."""
    assert _viewer_kwargs({})["beams"] is True
    assert _viewer_kwargs({"beams": False})["beams"] is False


def test_beam_endpoints_show_incident_and_reflected_rays():
    """Tests the beam arrows meet at the top of the rendered scene."""
    (incident_tail, hit), (reflected_start, reflected_tip) = _beam_endpoints(
        view_half=15.0, scene_z_max=8.0)

    assert np.allclose(hit, reflected_start)
    assert incident_tail[0] < hit[0] < reflected_tip[0]
    assert incident_tail[1] == hit[1] == reflected_tip[1] == 0.0
    assert hit[2] == 8.0
    assert incident_tail[2] > 8.0
    assert reflected_tip[2] > reflected_start[2]


def test_axis_label_color_contrasts_with_background():
    """Tests white and black backgrounds get opposite axis-label colors."""
    assert _contrasting_text_color((1.0, 1.0, 1.0)) == (0.0, 0.0, 0.0)
    assert _contrasting_text_color((0.0, 0.0, 0.0)) == (1.0, 1.0, 1.0)


def test_arrow_transform_keeps_absolute_radius_scale():
    """Tests arrow length scaling does not scale shaft and tip radii."""
    transform = _arrow_transform((1.0, 2.0, 3.0), (5.0, 2.0, 3.0))

    assert np.allclose(transform[:3, 3], (1.0, 2.0, 3.0))
    assert np.isclose(np.linalg.norm(transform[:3, 0]), 4.0)
    assert np.isclose(np.linalg.norm(transform[:3, 1]), 1.0)
    assert np.isclose(np.linalg.norm(transform[:3, 2]), 1.0)


if __name__ == "__main__":
    test_beams_are_enabled_by_default()
    test_beam_endpoints_show_incident_and_reflected_rays()
    test_axis_label_color_contrasts_with_background()
    test_arrow_transform_keeps_absolute_radius_scale()
