import pytest

import eradiate
from eradiate import unit_registry as ureg


@pytest.mark.parametrize("geometry", ["plane_parallel", "spherical_shell"])
def test_heterogeneous_parameter_lookup(
    modes_all_double, geometry, atmosphere_us_standard_mono, atmosphere_us_standard_ckd
):
    """
    Kernel parameter lookup
    =======================

    This test checks if parameters belonging to objects referenced by Mitsuba
    ``Medium`` objects are correctly retrieved.

    Rationale
    ---------

    Parameter paths are generated upon Mitsuba scene traversal based on the
    order in which they are discovered. The traversal order, although
    deterministic, is not easy to predict. Consequently, Eradiate implements a
    parameter lookup system which seeks a parameter during Mitsuba scene
    traversal based on its ID. This allows to safely retrieve Mitsuba parameter
    paths, regardless the number of objects which can register them as children
    or their discovery order.

    For that purpose, we create an ``Experiment`` and populate it with an
    atmosphere and an *in situ* measure. This way, the ``Medium`` instance
    generated by the atmosphere configuration is referenced both by the
    atmosphere shape and the sensor.

    Expected behaviour
    ------------------

    The sensor is discovered first because of the ID it is assigned (``dflux``).
    Consequently, although the medium is emitted upon atmosphere processing, the
    parameter associated with the volume data buffer nested in the ``Medium`` has
    the sensor in its parent nodes rather than the atmosphere shape.
    """
    atmosphere_mono = atmosphere_us_standard_mono
    atmosphere_mono.update({"has_scattering": False})

    atmosphere_ckd = atmosphere_us_standard_ckd
    atmosphere_ckd.update({"has_scattering": False})

    exp = eradiate.experiments.AtmosphereExperiment(
        geometry=geometry,
        surface={"type": "lambertian", "reflectance": 1.0},
        atmosphere=atmosphere_ckd if eradiate.mode().is_ckd else atmosphere_mono,
        illumination={"type": "directional", "zenith": 30.0},
        measures={
            "type": "distant_flux",
            "id": "dflux",
            "srf": {
                "type": "multi_delta",
                "wavelengths": 550.0 * ureg.nm,
            },
            "ray_offset": 1.0 * ureg.m,
        },
    )
    exp.init()

    # Medium is resolved, regardless the fact that it is first encountered as
    # a member of the "dflux" sensor
    assert (
        exp.mi_scene.umap_template["medium_atmosphere.albedo.volume.data"].parameter_id
        == "dflux.medium.albedo.volume.data"
        if geometry == "spherical_shell"
        else "dflux.medium.albedo.data"
    )
