# Path: dataflow/illumination_integration/helper.py
# def getBilinearFromUv(env_map: tf.Tensor, uvs: tf.Tensor) -> tf.Tensor:
#     u = uvs[..., 0]
#     v = uvs[..., 1]
# 
#     # vFlipped = 1 - v
# 
#     # u corresponds to width and v to heights
#     u_reshaped = tf.reshape(u, (-1, tf.math.reduce_prod(u.shape[1:])))
#     v_reshaped = tf.reshape(v, (-1, tf.math.reduce_prod(v.shape[1:])))
# 
#     uvs = tf.stack(
#         [u_reshaped * (env_map.shape[2] - 1), v_reshaped * (env_map.shape[1] - 1)],
#         axis=-1,
#     )
# 
#     return tfa.image.interpolate_bilinear(env_map, uvs, indexing="xy")
# 
# Path: nn_utils/math_utils.py
# def shape_to_uv(height: int, width: int) -> tf.Tensor:
#     # UV
#     # 0,0              1,0
#     # 0,1              1,1
#     us, vs = tf.meshgrid(
#         tf.linspace(
#             0.0 + 0.5 / tf.cast(width, tf.float32),
#             1.0 - 0.5 / tf.cast(width, tf.float32),
#             width,
#         ),
#         tf.linspace(
#             0.0 + 0.5 / tf.cast(height, tf.float32),
#             1.0 - 0.5 / tf.cast(height, tf.float32),
#             height,
#         ),
#     )  # Use pixel centers
#     return tf.cast(tf.stack([us, vs], -1), tf.float32)
# 
# def uv_to_direction(uvs: tf.Tensor) -> tf.Tensor:
#     spherical = uv_to_spherical(uvs)
#     theta = spherical[..., 0:1]
#     phi = spherical[..., 1:2]
#     return spherical_to_cartesian(theta, phi)


from typing import List, Union
from dataflow.illumination_integration.helper import (
    getBilinearFromUv,
)
from nn_utils.math_utils import shape_to_uv, uv_to_direction
import numpy as np
import tensorflow as tf
import nn_utils.math_utils as math_utils
    total_directions_required = num_roughness_0 + num_random_roughness

    if uvs_flat.shape[0] < total_directions_required:
        repeats_required = tf.cast(
            tf.math.ceil(total_directions_required / uvs_flat.shape[0]), tf.int32
        )
        uvs_flat = math_utils.repeat(uvs_flat, repeats_required, 0)

    uvs_shuffle = tf.random.shuffle(uvs_flat)
    uvs_random = uvs_shuffle[:total_directions_required]

    jitter = tf.random.normal(uvs_random.shape, mean=0.0, stddev=0.3)
    uvs_random = uvs_random + jitter

    # Setup roughness
    roughness_random = tf.clip_by_value(
        tf.random.uniform(
            (num_random_roughness, 1), minval=1 / 255, maxval=1 + 1 / 255
        ),
        0,
        1,
    )

    r0_uvs = uvs_random[:num_roughness_0]
    rnd_uvs = uvs_random[num_roughness_0 : num_roughness_0 + num_random_roughness]

    # Get samples
    samples_random = random_uv_roughness_access(data_levels, rnd_uvs, roughness_random)

    # Always get r0 samples

# Path: dataflow/illumination_integration/helper.py
# def getBilinearFromUv(env_map: tf.Tensor, uvs: tf.Tensor) -> tf.Tensor:
#     u = uvs[..., 0]
#     v = uvs[..., 1]
# 
#     # vFlipped = 1 - v
# 
#     # u corresponds to width and v to heights
#     u_reshaped = tf.reshape(u, (-1, tf.math.reduce_prod(u.shape[1:])))
#     v_reshaped = tf.reshape(v, (-1, tf.math.reduce_prod(v.shape[1:])))
# 
#     uvs = tf.stack(
#         [u_reshaped * (env_map.shape[2] - 1), v_reshaped * (env_map.shape[1] - 1)],
#         axis=-1,
#     )
# 
#     return tfa.image.interpolate_bilinear(env_map, uvs, indexing="xy")
# 
# Path: nn_utils/math_utils.py
# def shape_to_uv(height: int, width: int) -> tf.Tensor:
#     # UV
#     # 0,0              1,0
#     # 0,1              1,1
#     us, vs = tf.meshgrid(
#         tf.linspace(
#             0.0 + 0.5 / tf.cast(width, tf.float32),
#             1.0 - 0.5 / tf.cast(width, tf.float32),
#             width,
#         ),
#         tf.linspace(
#             0.0 + 0.5 / tf.cast(height, tf.float32),
#             1.0 - 0.5 / tf.cast(height, tf.float32),
#             height,
#         ),
#     )  # Use pixel centers
#     return tf.cast(tf.stack([us, vs], -1), tf.float32)
# 
# def uv_to_direction(uvs: tf.Tensor) -> tf.Tensor:
#     spherical = uv_to_spherical(uvs)
#     theta = spherical[..., 0:1]
#     phi = spherical[..., 1:2]
#     return spherical_to_cartesian(theta, phi)


# Path: dataflow/illumination_integration/dataflow.py
from typing import List, Union
from dataflow.illumination_integration.helper import (
    getBilinearFromUv,
)
from nn_utils.math_utils import shape_to_uv, uv_to_direction
import numpy as np
import tensorflow as tf
import nn_utils.math_utils as math_utils

    total_directions_required = num_roughness_0 + num_random_roughness

    if uvs_flat.shape[0] < total_directions_required:
        repeats_required = tf.cast(
            tf.math.ceil(total_directions_required / uvs_flat.shape[0]), tf.int32
        )
        uvs_flat = math_utils.repeat(uvs_flat, repeats_required, 0)

    uvs_shuffle = tf.random.shuffle(uvs_flat)
    uvs_random = uvs_shuffle[:total_directions_required]

    jitter = tf.random.normal(uvs_random.shape, mean=0.0, stddev=0.3)
    uvs_random = uvs_random + jitter

    # Setup roughness
    roughness_random = tf.clip_by_value(
        tf.random.uniform(
            (num_random_roughness, 1), minval=1 / 255, maxval=1 + 1 / 255
        ),
        0,
        1,
    )

    r0_uvs = uvs_random[:num_roughness_0]
    rnd_uvs = uvs_random[num_roughness_0 : num_roughness_0 + num_random_roughness]

    # Get samples
    samples_random = random_uv_roughness_access(data_levels, rnd_uvs, roughness_random)

    # Always get r0 samples

    samples_r0 = random_uv_roughness_access(