"""A series of basic one-dimensional test cases."""

from contextlib import nullcontext

import mitsuba as mi
import numpy as np
import pytest

import eradiate
from eradiate import KernelContext
from eradiate import unit_registry as ureg
from eradiate.kernel import mi_render
from eradiate.scenes.bsdfs import LambertianBSDF
from eradiate.scenes.core import Scene
from eradiate.scenes.illumination import ConstantIllumination, DirectionalIllumination, AstroObjectIllumination
from eradiate.scenes.integrators import PathIntegrator
from eradiate.scenes.measure import MultiDistantMeasure
from eradiate.scenes.shapes import RectangleShape
from eradiate.scenes.surface import BasicSurface
from eradiate.test_tools.types import check_scene_element


@pytest.mark.parametrize("illumination, spp", [("directional", 1), ("constant", 5e5), ("astroobject", 5e5)])
@pytest.mark.parametrize("li", [0.1, 1.0, 10.0])
@pytest.mark.slow
def test_radiometric_accuracy(modes_all_mono, illumination, spp, li, ert_seed_state):
    r"""
    Radiometric check (``path``)
    ============================

    This simple test case compares simulated leaving radiance at a Lambertian
    surface with theoretical values.

    Rationale
    ---------

    * Geometry: a square surface with unit size and a Lambertian BRDF with
      reflectance :math:`\rho = 0.5`.
    * Illumination:

      * ``directional``: a directional light source at the zenith with
        radiance :math:`L_\mathrm{i} \in [0.1, 1, 10]`.
      * ``constant``: an isotropic illumination with radiance
        :math:`L_\mathrm{i} \in [0.1, 1, 10]`.

    * Sensor: A ``distant`` sensor targeting (0, 0, 0).


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

    The computed solution is equal to the theoretical solution (relative
    tolerance of 0.1%).

    Theoretical solutions:

    * ``directional``: :math:`L_\mathrm{o} = \frac{\rho L_\mathrm{i}}{\pi}`
    * ``constant``: :math:`L_\mathrm{o} = \rho L_\mathrm{i}`
    """

    # Basic configuration
    vza = np.linspace(0, 80, 10)
    rho = 0.5

    # Ignore warning in single precision
    with pytest.warns(UserWarning, match="the selected mode is single-precision") if (
        eradiate.mode().is_single_precision and spp > 100000
    ) else nullcontext():
        measure = MultiDistantMeasure.hplane(
            zeniths=vza, azimuth=0.0, target=[0, 0, 0], spp=spp
        )

    objects = {
        "surface": BasicSurface(
            bsdf=LambertianBSDF(reflectance=rho),
            shape=RectangleShape(edges=2.0 * ureg.m),
        ),
        "measure": measure,
        "integrator": PathIntegrator(),
    }

    if illumination == "directional":
        objects["illumination"] = DirectionalIllumination(zenith=0.0, irradiance=li)
        theoretical_solution = np.full_like(vza, rho * li / np.pi)
        rtol=1e-3

    elif illumination == "constant":
        objects["illumination"] = ConstantIllumination(radiance=li)
        theoretical_solution = np.full_like(vza, rho * li)
        rtol=1e-3

    elif illumination == "astroobject":
        objects["illumination"] = AstroObjectIllumination(zenith=0.0, irradiance=li, angular_diameter=0.03)
        theoretical_solution = np.full_like(vza, rho * li / np.pi)
        # The angular diameter is not taken into account in the theoretical solution
        rtol=1e-2

    else:
        raise ValueError(f"unsupported illumination '{illumination}'")

    scene = Scene(objects=objects)
    mi_wrapper = check_scene_element(scene, mi.Scene)

    result = np.squeeze(mi_render(mi_wrapper, ctxs=[KernelContext()])[550.0]["measure"])
    np.testing.assert_allclose(result, theoretical_solution, rtol=rtol)
