from pathlib import Path

import chex
import equinox as eqx
import jax.numpy as jnp
import numpy as np
import pytest
from pytest_subtests import SubTests

from differt.em import materials
from differt.geometry import (
    Mesh,
    Scene,
    assemble_path,
    fibonacci_lattice,
    first_triangle_hit_by_ray,
    path_length,
    ray_intersect_any_triangle,
    ray_intersect_triangle,
)
from differt_core.geometry import SionnaScene


@pytest.mark.slow
def test_ray_casting() -> None:
    o3d = pytest.importorskip("open3d", reason="open3d not installed")

    knot_mesh = o3d.data.KnotMesh()
    o3d_mesh = o3d.t.io.read_triangle_mesh(knot_mesh.path).translate([50, 20, 10])
    o3d_mesh = o3d_mesh.compute_vertex_normals()  # This avoids a warning from Open3D
    o3d_mesh = o3d_mesh.compute_triangle_normals()

    mesh = Mesh(
        vertices=jnp.asarray(o3d_mesh.vertex.positions.numpy()),
        triangles=jnp.asarray(o3d_mesh.triangle.indices.numpy()),
    )

    chex.assert_trees_all_close(
        mesh.bounding_box,
        np.stack(
            [
                o3d_mesh.get_min_bound().numpy(),
                o3d_mesh.get_max_bound().numpy(),
            ],
            axis=0,
        ),
    )

    chex.assert_trees_all_close(
        mesh.normals, o3d_mesh.triangle.normals.numpy(), atol=1e-6
    )

    scene = o3d.t.geometry.RaycastingScene()
    scene.add_triangles(o3d_mesh)

    ray_directions = fibonacci_lattice(1_000)
    ray_directions = fibonacci_lattice(50)
    ray_origins = jnp.zeros_like(ray_directions)

    o3d_rays = o3d.core.Tensor(
        np.concatenate((ray_origins, ray_directions), axis=-1),
        dtype=o3d.core.Dtype.Float32,
    )

    triangle_vertices = mesh.triangle_vertices

    triangles, t_hit = first_triangle_hit_by_ray(
        ray_origins, ray_directions, triangle_vertices
    )
    hit = triangles != -1
    triangles = triangles.astype(jnp.uint32)

    ans = scene.cast_rays(o3d_rays, nthreads=1)  # codespell:ignore ans

    chex.assert_trees_all_close(
        t_hit,
        ans["t_hit"].numpy(),  # codespell:ignore ans
        atol=1e-4,
    )
    chex.assert_trees_all_equal(
        jnp.where(hit, triangles, jnp.asarray(scene.INVALID_ID, dtype=jnp.uint32)),
        ans["primitive_ids"].numpy(),  # codespell:ignore ans
    )

    got_counts = ray_intersect_triangle(
        ray_origins[..., None, :], ray_directions[..., None, :], triangle_vertices
    )[1].sum(axis=-1)

    expected_counts = scene.count_intersections(o3d_rays, nthreads=1).numpy()

    chex.assert_trees_all_equal(
        got_counts,
        expected_counts,
    )

    scale = 100.0

    got_hit = ray_intersect_any_triangle(
        ray_origins,
        scale * ray_directions,
        triangle_vertices,
    )

    expected_hit = scene.test_occlusions(o3d_rays, tfar=scale, nthreads=1).numpy()

    chex.assert_trees_all_equal(
        got_hit,
        expected_hit,
    )


@pytest.mark.slow
def test_simple_street_canyon() -> None:
    mi = pytest.importorskip("mitsuba", reason="mitsuba not installed")
    try:
        mi.set_variant("llvm_ad_mono_polarized")
    except AttributeError:
        pytest.skip("Mitsuba variant 'llvm_ad_mono_polarized' not available")
    sionna = pytest.importorskip("sionna", reason="sionna not installed")
    file = sionna.rt.scene.simple_street_canyon

    sionna_scene = sionna.rt.load_scene(file)
    differt_scene = Scene.load_xml(file).set_assume_quads()  # Faster RT

    sionna_scene.tx_array = sionna.rt.PlanarArray(
        num_rows=1,
        num_cols=1,
        vertical_spacing=0.5,
        horizontal_spacing=0.5,
        pattern="tr38901",
        polarization="V",
    )

    sionna_scene.rx_array = sionna.rt.PlanarArray(
        num_rows=1,
        num_cols=1,
        vertical_spacing=0.5,
        horizontal_spacing=0.5,
        pattern="dipole",
        polarization="cross",
    )

    tx = sionna.rt.Transmitter(name="tx", position=[-33.0, 0.0, 32.0])

    sionna_scene.add(tx)

    rx = sionna.rt.Receiver(name="rx", position=[20.0, 0.0, 2.0], orientation=[0, 0, 0])

    sionna_scene.add(rx)

    tx.look_at(rx)

    differt_scene = eqx.tree_at(
        lambda s: s.transmitters,
        differt_scene,
        replace=tx.position.jax().reshape(3),
    )

    differt_scene = eqx.tree_at(
        lambda s: s.receivers,
        differt_scene,
        replace=rx.position.jax().reshape(3),
    )

    max_order = 4

    sionna_solver = sionna.rt.PathSolver()
    sionna_paths = sionna_solver(sionna_scene, max_depth=max_order, refraction=False)
    sionna_path_objects = sionna_paths.objects.jax()
    sionna_path_vertices = sionna_paths.vertices.jax()

    max_depth = sionna_path_objects.shape[0]  # May differ from 'max_order'

    for order in range(max_depth + 1):
        paths = differt_scene.trace_paths(
            order=order,
            solver="hybrid",
        )
        select = (sionna_path_objects == -1).sum(axis=0) == (max_depth - order)
        vertices = sionna_path_vertices[:order, select, :]
        vertices = jnp.moveaxis(vertices, 0, -2)
        vertices = assemble_path(
            differt_scene.transmitters,
            vertices,
            differt_scene.receivers,
        )
        got_path_lengths = path_length(paths.masked_vertices)
        expected_path_lengths = path_length(vertices)
        # We check the sum of path lengths because Sionna orders the paths differently,
        # so we cannot compare them directly.
        chex.assert_trees_all_close(
            got_path_lengths.sum(),
            expected_path_lengths.sum(),
            atol=1e-5,
            custom_message=f"Mismatch for paths {order = }, differt = {paths.masked_vertices!r}, sionna = {vertices!r}.",
        )


# Sionna RT still ships a pre-P.2040-4 ITU materials list: some materials are
# missing entirely, others have coefficients that don't match DiffeRT's. An
# upstream PR updates Sionna RT's list, but hasn't been merged (nor released)
# yet, and DiffeRT's own coefficients for a couple of materials don't match that
# PR either (possibly transcription discrepancies worth checking against the
# primary ITU-R P.2040-4 recommendation):
_SIONNA_ITU_UPDATE_PR = "https://github.com/NVlabs/sionna-rt/pull/70"
_MATERIALS_MISSING_FROM_SIONNA = frozenset({
    "itu_vacuum",
    "itu_clear_acrylic",
    "itu_vinyl_tile",
    "itu_carpet_tile",
    "itu_asphalt_concrete",
})
_MATERIALS_WITH_MISMATCHING_COEFFICIENTS = frozenset({
    "itu_glass",
    "itu_ceiling_board",
    "itu_plasterboard",
    "itu_brick",
})


def _differt_color(itu_type: str, tmp_path: Path) -> tuple[float, float, float]:
    """Parse the color DiffeRT assigns to an ITU radio material, via a minimal Sionna scene."""
    xml = f"""<scene version="2.1.0">
  <bsdf type="itu-radio-material" id="mat-0">
    <string name="type" value="{itu_type}"/>
  </bsdf>
  <shape type="ply" id="mesh-0">
    <string name="filename" value="dummy.ply"/>
    <ref id="mat-0"/>
  </shape>
</scene>
"""
    file = tmp_path / "scene.xml"
    file.write_text(xml)
    return tuple(SionnaScene.load_xml(file).materials["mat-0"].color)


def test_itu_materials(subtests: SubTests, tmp_path: Path) -> None:
    mi = pytest.importorskip("mitsuba", reason="mitsuba not installed")
    try:
        mi.set_variant("llvm_ad_mono_polarized")
    except AttributeError:
        pytest.skip("Mitsuba variant 'llvm_ad_mono_polarized' not available")
    sionna = pytest.importorskip("sionna", reason="sionna not installed")

    for differt_mat in materials.values():
        # `materials` maps both official ITU names (e.g., "Wood") and Sionna-style
        # aliases (e.g., "itu_wood") to the same `Material` instance, but iterating
        # only sees the primary (official) name, so we look up the alias instead.
        mat_name = next((a for a in differt_mat.aliases if a.startswith("itu_")), None)
        if mat_name is None:
            continue

        itu_type = mat_name.removeprefix("itu_")
        missing_from_sionna = mat_name in _MATERIALS_MISSING_FROM_SIONNA
        mismatching_coefficients = mat_name in _MATERIALS_WITH_MISMATCHING_COEFFICIENTS

        with subtests.test(f"{mat_name} color"):
            try:
                got_color = _differt_color(itu_type, tmp_path)
                expected_color = sionna.rt.ITURadioMaterial.ITU_MATERIAL_COLORS[
                    itu_type
                ]
                chex.assert_trees_all_close(got_color, expected_color, atol=1e-6)
            except (KeyError, AssertionError) as exc:
                if not missing_from_sionna:
                    raise
                pytest.xfail(
                    f"Sionna RT doesn't define a color for {itu_type!r} yet, see "
                    f"{_SIONNA_ITU_UPDATE_PR} ({exc})."
                )
            else:
                if missing_from_sionna:
                    pytest.fail(
                        f"Sionna RT now defines a color for {itu_type!r}; remove it "
                        f"from _MATERIALS_MISSING_FROM_SIONNA ({_SIONNA_ITU_UPDATE_PR})."
                    )

        if missing_from_sionna or mismatching_coefficients:
            with subtests.test(
                f"{mat_name} properties (tracking {_SIONNA_ITU_UPDATE_PR})"
            ):
                f = 10e9
                try:
                    sionna_mat_relative_permittivity, sionna_mat_conductivity = (
                        sionna.rt.radio_materials.itu.itu_material(itu_type, f)
                    )
                    chex.assert_trees_all_close(
                        differt_mat.relative_permittivity(f),
                        sionna_mat_relative_permittivity,
                    )
                    chex.assert_trees_all_close(
                        differt_mat.conductivity(f), sionna_mat_conductivity
                    )
                except (ValueError, AssertionError) as exc:
                    pytest.xfail(
                        f"Sionna RT's coefficients for {itu_type!r} don't match "
                        f"DiffeRT's yet, see {_SIONNA_ITU_UPDATE_PR} ({exc})."
                    )
                else:
                    pytest.fail(
                        f"Sionna RT's coefficients for {itu_type!r} now match "
                        f"DiffeRT's; remove it from "
                        f"_MATERIALS_WITH_MISMATCHING_COEFFICIENTS "
                        f"({_SIONNA_ITU_UPDATE_PR})."
                    )
            continue

        # We multiply by 1.1 to avoid checking on freq. limits, because Sionna will fail
        for f in 1.1 * np.logspace(9 - 2, 9 + 3, 21):
            differt_mat_relative_permittivity = differt_mat.relative_permittivity(f)
            differt_mat_conductivity = differt_mat.conductivity(f)
            differt_out_of_range = (
                differt_mat_relative_permittivity,
                differt_mat_conductivity,
            ) == (-1.0, -1.0)

            with subtests.test(f"{mat_name} @ {f / 1e9} GHz"):
                try:
                    sionna_mat_relative_permittivity, sionna_mat_conductivity = (
                        sionna.rt.radio_materials.itu.itu_material(itu_type, f)
                    )
                except ValueError as exc:
                    if differt_out_of_range:
                        # Both sides agree this frequency is out of range.
                        continue
                    # DiffeRT's P.2040-4 update extends the frequency range for some
                    # materials beyond what Sionna RT's (older) table covers.
                    pytest.xfail(
                        f"Sionna RT doesn't cover {f / 1e9:.3g} GHz for {itu_type!r} "
                        f"yet, see {_SIONNA_ITU_UPDATE_PR} ({exc})."
                    )
                else:
                    if differt_out_of_range:
                        pytest.fail(
                            f"DiffeRT considers {itu_type!r} undefined at "
                            f"{f / 1e9:.3g} GHz, but Sionna RT returned a value; "
                            "DiffeRT's frequency-range logic may need updating."
                        )

                    chex.assert_trees_all_close(
                        differt_mat_relative_permittivity,
                        sionna_mat_relative_permittivity,
                        custom_message=f"Mismatch for {mat_name = } @ {f / 1e9} GHz.",
                    )

                    chex.assert_trees_all_close(
                        differt_mat_conductivity,
                        sionna_mat_conductivity,
                        custom_message=f"Mismatch for {mat_name = } @ {f / 1e9} GHz.",
                    )
