import asyncio
import time
from collections import defaultdict
import pathlib
from PIL import Image
from PIL import ImageEnhance
import random
import omni.kit.test_helpers_gfx.compare_utils
import carb
import carb.tokens
import omni.kit.app
import omni.kit.commands
import omni.kit.test
import omni.ui as ui
import omni.usd.audio
import os
from omni.ui.tests.test_base import OmniUiTest
from omni.kit.test_suite.helpers import wait_stage_loading


OUTPUTS_DIR = pathlib.Path(omni.kit.test.get_test_output_path())

def wait_for_asset(test, audio, prim):
    i = 0
    while audio.get_sound_asset_status(prim) == omni.usd.audio.AssetLoadStatus.IN_PROGRESS:
        time.sleep(0.001)
        if i > 5000:
            raise Exception("asset load timed out")
        i += 1

    test.assertEqual(audio.get_sound_asset_status(prim), omni.usd.audio.AssetLoadStatus.DONE)

async def wait_for_prim_change(test, context, audio, prim):
    await context.next_frame_async()

    # the asset may reload, so we need to wait for it
    wait_for_asset(test, audio, prim)

class TestStageAudioMinimal(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        self._context = omni.usd.get_context()
        self.assertIsNotNone(self._context)

        self._context.new_stage()
        self._stage = self._context.get_stage()
        self.assertIsNotNone(self._stage)

        self._audio = omni.usd.audio.get_stage_audio_interface()
        self.assertIsNotNone(self._audio)

    async def test_hydra_plugin(self):
        # check that the hydra plugin still works
        self.assertTrue(omni.usd.audio.test_hydra_plugin(), "the hydra audio plugin failed to load")

    async def test_device(self):
        self._audio.set_device("default")
        time.sleep(2.0)

        self._audio.set_device("apple")
        time.sleep(2.0)

        self._audio.set_device("Speakers")
        time.sleep(2.0)

    async def test_set_doppler_default(self):
        self._audio.set_doppler_default(omni.usd.audio.FeatureDefault.FORCE_OFF)
        self.assertEqual(self._audio.get_doppler_default(), omni.usd.audio.FeatureDefault.FORCE_OFF)

        self._audio.set_doppler_default()
        self.assertEqual(self._audio.get_doppler_default(), omni.usd.audio.FeatureDefault.OFF)

        # FIXME: This is just verifying that we can get/set this value.
        #        This test should somehow verify that this feature is actually
        #        detectable in the output audio.

    async def test_set_distance_delay_default(self):
        self._audio.set_distance_delay_default(omni.usd.audio.FeatureDefault.FORCE_OFF)
        self.assertEqual(self._audio.get_distance_delay_default(), omni.usd.audio.FeatureDefault.FORCE_OFF)

        self._audio.set_distance_delay_default()
        self.assertEqual(self._audio.get_distance_delay_default(), omni.usd.audio.FeatureDefault.OFF)

        # FIXME: This is just verifying that we can get/set this value.
        #        This test should somehow verify that this feature is actually
        #        detectable in the output audio.

    async def test_set_interaural_delay_default(self):
        self._audio.set_interaural_delay_default(omni.usd.audio.FeatureDefault.FORCE_OFF)
        self.assertEqual(self._audio.get_interaural_delay_default(), omni.usd.audio.FeatureDefault.FORCE_OFF)

        self._audio.set_interaural_delay_default()
        self.assertEqual(self._audio.get_interaural_delay_default(), omni.usd.audio.FeatureDefault.OFF)

        # FIXME: This is just verifying that we can get/set this value.
        #        This test should somehow verify that this feature is actually
        #        detectable in the output audio.

    async def test_set_concurrent_voices(self):
        self._audio.set_concurrent_voices(32)
        self.assertEqual(self._audio.get_concurrent_voices(), 32)

        self._audio.set_concurrent_voices()
        self.assertEqual(self._audio.get_concurrent_voices(), 64)

        # FIXME: This is just verifying that we can get/set this value.
        #        This test should somehow verify that this feature is actually
        #        detectable in the output audio.

    async def test_speed_of_sound(self):
        self._audio.set_speed_of_sound(120.0)
        self.assertEqual(self._audio.get_speed_of_sound(), 120.0)

        self._audio.set_speed_of_sound()
        self.assertLess(abs(self._audio.get_speed_of_sound() - 340.0), 0.1)

        # FIXME: This is just verifying that we can get/set this value.
        #        This test should somehow verify that this feature is actually
        #        detectable in the output audio.

    async def test_doppler_scale(self):
        self._audio.set_doppler_scale(4.0)
        self.assertEqual(self._audio.get_doppler_scale(), 4.0)

        self._audio.set_doppler_scale()
        self.assertEqual(self._audio.get_doppler_scale(), 1.0)

        # FIXME: This is just verifying that we can get/set this value.
        #        This test should somehow verify that this feature is actually
        #        detectable in the output audio.

    async def test_doppler_limit(self):
        self._audio.set_doppler_limit(4.0)
        self.assertEqual(self._audio.get_doppler_limit(), 4.0)

        self._audio.set_doppler_limit()
        self.assertEqual(self._audio.get_doppler_limit(), 2.0)

        # FIXME: This is just verifying that we can get/set this value.
        #        This test should somehow verify that this feature is actually
        #        detectable in the output audio.

    async def test_spatial_time_scale(self):
        self._audio.set_spatial_time_scale(4.0)
        self.assertEqual(self._audio.get_spatial_time_scale(), 4.0)

        self._audio.set_spatial_time_scale()
        self.assertEqual(self._audio.get_spatial_time_scale(), 1.0)

        # FIXME: This is just verifying that we can get/set this value.
        #        This test should somehow verify that this feature is actually
        #        detectable in the output audio.

    async def test_nonspatial_time_scale(self):
        self._audio.set_nonspatial_time_scale(4.0)
        self.assertEqual(self._audio.get_nonspatial_time_scale(), 4.0)

        self._audio.set_nonspatial_time_scale()
        self.assertEqual(self._audio.get_nonspatial_time_scale(), 1.0)

        # FIXME: This is just verifying that we can get/set this value.
        #        This test should somehow verify that this feature is actually
        #        detectable in the output audio.

    async def test_get_metadata_change_stream(self):
        callbacks = 0
        event_type = None
        def event_callback(event):
            nonlocal event_type
            nonlocal callbacks
            event_type = event.type
            callbacks += 1

        def wait_for_callback(initial):
            nonlocal callbacks
            i = 0
            while callbacks == initial:
                time.sleep(0.001)
                if i > 1000:
                    raise Exception("asset load timed out")
                i += 1

        self._events = self._audio.get_metadata_change_stream()
        self.assertIsNotNone(self._events)

        self._stage_event_sub = self._events.create_subscription_to_pop(
            event_callback, name="stage audio test callback"
        )

        initial = callbacks
        self._audio.set_doppler_default(omni.usd.audio.FeatureDefault.FORCE_OFF)
        wait_for_callback(initial)
        self.assertEqual(callbacks, initial + 1)
        self.assertEqual(event_type, int(omni.usd.audio.EventType.METADATA_CHANGE))

        # try the same thing twice to verify that the value needs to change to
        # trigger a callback
        initial = callbacks
        self._audio.set_doppler_default(omni.usd.audio.FeatureDefault.FORCE_OFF)
        try:
            wait_for_callback(initial)
            self.assertTrue(False)
        except Exception:
            pass

        self.assertEqual(callbacks, initial)

        initial = callbacks
        self._audio.set_distance_delay_default(omni.usd.audio.FeatureDefault.FORCE_OFF)
        wait_for_callback(initial)
        self.assertEqual(callbacks, initial + 1)
        self.assertEqual(event_type, int(omni.usd.audio.EventType.METADATA_CHANGE))

        initial = callbacks
        self._audio.set_interaural_delay_default(omni.usd.audio.FeatureDefault.FORCE_OFF)
        wait_for_callback(initial)
        self.assertEqual(callbacks, initial + 1)
        self.assertEqual(event_type, int(omni.usd.audio.EventType.METADATA_CHANGE))

        initial = callbacks
        self._audio.set_concurrent_voices(32)
        wait_for_callback(initial)
        self.assertEqual(callbacks, initial + 1)
        self.assertEqual(event_type, int(omni.usd.audio.EventType.METADATA_CHANGE))

        initial = callbacks
        self._audio.set_speed_of_sound(120.0)
        wait_for_callback(initial)
        self.assertEqual(callbacks, initial + 1)
        self.assertEqual(event_type, int(omni.usd.audio.EventType.METADATA_CHANGE))

        initial = callbacks
        self._audio.set_doppler_scale(4.0)
        wait_for_callback(initial)
        self.assertEqual(callbacks, initial + 1)
        self.assertEqual(event_type, int(omni.usd.audio.EventType.METADATA_CHANGE))

        initial = callbacks
        self._audio.set_doppler_limit(4.0)
        wait_for_callback(initial)
        self.assertEqual(callbacks, initial + 1)
        self.assertEqual(event_type, int(omni.usd.audio.EventType.METADATA_CHANGE))

        initial = callbacks
        self._audio.set_spatial_time_scale(4.0)
        wait_for_callback(initial)
        self.assertEqual(callbacks, initial + 1)
        self.assertEqual(event_type, int(omni.usd.audio.EventType.METADATA_CHANGE))

        initial = callbacks
        self._audio.set_nonspatial_time_scale(4.0)
        wait_for_callback(initial)
        self.assertEqual(callbacks, initial + 1)
        self.assertEqual(event_type, int(omni.usd.audio.EventType.METADATA_CHANGE))


class TestStageAudio(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        self._context = omni.usd.get_context()
        self.assertIsNotNone(self._context)

        self._context.new_stage()
        self._stage = self._context.get_stage()
        self.assertIsNotNone(self._stage)

        self._audio = omni.usd.audio.get_stage_audio_interface()
        self.assertIsNotNone(self._audio)

        extension_path = carb.tokens.get_tokens_interface().resolve("${omni.kit.audio.test.usd}")
        self._test_path = pathlib.Path(extension_path).joinpath("data").joinpath("tests").absolute()
        self._golden_path = self._test_path.joinpath("golden")

        self._prim_no_filename_path = "/sound_no_filename"
        self._prim_no_filename = self._stage.DefinePrim(self._prim_no_filename_path, "Sound")
        self.assertIsNotNone(self._prim_no_filename)

        self._prim_bad_asset_path = "/sound_bad_asset"
        self._prim_bad_asset = self._stage.DefinePrim(self._prim_bad_asset_path, "Sound")
        self.assertIsNotNone(self._prim_bad_asset)
        self._prim_bad_asset.GetAttribute("filePath").Set("/a/b/c/d/e/f/g/h/i/j/k/l/m/n/o/p/q/r/s/t/u/v/x/y/z")

        self._test_sound_1ch_path = self._test_path.joinpath("short-1ch.oga").absolute()
        self._prim_1ch_path = "/sound_1ch"
        self._prim_1ch = self._stage.DefinePrim(self._prim_1ch_path, "Sound")
        self.assertIsNotNone(self._prim_1ch)
        self._prim_1ch.GetAttribute("filePath").Set(str(self._test_sound_1ch_path))

        self._test_sound_2ch_path = self._test_path.joinpath("short-2ch.oga").absolute()
        self._prim_2ch_path = "/sound_2ch"
        self._prim_2ch = self._stage.DefinePrim(self._prim_2ch_path, "Sound")
        self.assertIsNotNone(self._prim_2ch)
        self._prim_2ch.GetAttribute("filePath").Set(str(self._test_sound_2ch_path))

        self._test_sound_long_path = self._test_path.joinpath("long.oga").absolute()
        self._prim_long_path = "/sound_long"
        self._prim_long = self._stage.DefinePrim(self._prim_long_path, "Sound")
        self.assertIsNotNone(self._prim_long)
        self._prim_long.GetAttribute("filePath").Set(str(self._test_sound_long_path))

        wait_for_asset(self, self._audio, self._prim_1ch)
        wait_for_asset(self, self._audio, self._prim_2ch)
        wait_for_asset(self, self._audio, self._prim_long)

        try:
            os.mkdir(OUTPUTS_DIR)
        except FileExistsError:
            pass


    async def tearDown(self):
        self._stage.RemovePrim(self._prim_no_filename_path)
        self._stage.RemovePrim(self._prim_bad_asset_path)
        self._stage.RemovePrim(self._prim_1ch_path)
        self._stage.RemovePrim(self._prim_2ch_path)
        self._stage.RemovePrim(self._prim_long_path)

    async def test_get_sound_count(self):
        # sounds are only counted if they have a valid asset
        self.assertEqual(self._audio.get_sound_count(), 3);

        self._audio.draw_waveform(self._prim_no_filename, 1, 1)
        self._audio.draw_waveform(self._prim_bad_asset, 1, 1)
        self.assertEqual(self._audio.get_sound_count(), 5);

    async def test_is_sound_playing(self):
        self.assertFalse(self._audio.is_sound_playing(self._prim_long))

        # the sound asset for this prim is silent, so the test shouldn't bother testers
        # play the sound
        self._audio.play_sound(self._prim_long)
        self.assertTrue(self._audio.is_sound_playing(self._prim_long))

        # stop the sound
        self._audio.stop_sound(self._prim_long)
        self.assertFalse(self._audio.is_sound_playing(self._prim_long))

        # do it again, but use stop_all_sounds() instead
        self._audio.play_sound(self._prim_long)
        self.assertTrue(self._audio.is_sound_playing(self._prim_long))
        self._audio.stop_all_sounds()
        self.assertFalse(self._audio.is_sound_playing(self._prim_long))

        # do it again but with spawn_voice()
        voice = self._audio.spawn_voice(self._prim_long)

        # the voice is not managed by IStageAudio
        self.assertFalse(self._audio.is_sound_playing(self._prim_long))
        self.assertTrue(voice.is_playing())
        voice.stop()
        self.assertFalse(voice.is_playing())

    async def test_subscribe_to_asset_load(self):
        loaded = False
        def load_callback():
            self.assertEqual(self._audio.get_sound_asset_status(self._prim_long), omni.usd.audio.AssetLoadStatus.DONE)
            nonlocal loaded
            loaded = True

        # try it with something that's already loaded to ensure a callback will occur in this case
        self.assertTrue(self._audio.subscribe_to_asset_load(self._prim_long, load_callback))
        while not loaded:
            time.sleep(0.001)
            if i > 5000:
                raise Exception("asset load timed out")
            i += 1

        # switch the file path and try again, so a load will actually be needed
        self._prim_2ch.GetAttribute("filePath").Set(str(self._test_sound_2ch_path))
        loaded = False
        self.assertTrue(self._audio.subscribe_to_asset_load(self._prim_long, load_callback))
        while not loaded:
            time.sleep(0.001)
            if i > 5000:
                raise Exception("asset load timed out")
            i += 1

    async def test_draw_waveform(self):
        # toggle in case you want to regenerate waveforms
        GENERATE_GOLDEN_IMAGES = False

        if GENERATE_GOLDEN_IMAGES:
            print("!!!!! Golden images are being regenerated to '" + str(pathlib.Path(OUTPUTS_DIR)) + "'.")
            print("!!!!! Once regenerated, these files should be copied to '" + str(self._golden_path) + "'.")
            print("!!!!! Note: the `GENERATE_GOLDEN_IMAGES` variable being set to 'True' above should never be committed.")

        W = 256
        H = 256
        BASE_NAME = "test_stage_audio.draw_waveform"
        TEST_NAMES = [
                BASE_NAME + ".1ch.USE_LINES.SPLIT_CHANNELS.png",
                BASE_NAME + ".1ch.USE_LINES.SPLIT_CHANNELS.BACKGROUNDCOLOR.png",
                BASE_NAME + ".2ch.USE_LINES.SPLIT_CHANNELS.png",
                BASE_NAME + ".2ch.USE_LINES.SPLIT_CHANNELS.PARTIAL_COLOR.png",
                BASE_NAME + ".2ch.USE_LINES.SPLIT_CHANNELS.COLOR.png",
                BASE_NAME + ".1ch.USE_LINES.SPLIT_CHANNELS.MEDIA_OFFSET_START.png",
                BASE_NAME + ".1ch.USE_LINES.SPLIT_CHANNELS.MEDIA_OFFSET_START.MEDIA_OFFSET_END.png",
                BASE_NAME + ".1ch.png",
                BASE_NAME + ".1ch.ALPHA_BLEND.png",
                BASE_NAME + ".2ch.ch_2.png",
                BASE_NAME + ".2ch.ALPHA_BLEND.MULTI_CHANNEL.png",
                BASE_NAME + ".2ch.NOISE_COLOR.2bit.png",
        ]

        raw = self._audio.draw_waveform(self._prim_no_filename, W, H)
        self.assertEqual(len(raw), 0)

        raw = self._audio.draw_waveform(self._prim_bad_asset, W, H)
        self.assertEqual(len(raw), 0)

        # render the mono waveform
        raw = self._audio.draw_waveform(self._prim_1ch, W, H)
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[0])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[0]),
                    self._golden_path.joinpath(TEST_NAMES[0]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[0] + ".diff.png")),
                0.1)

        # render it with a custom background color
        raw = self._audio.draw_waveform(self._prim_1ch, W, H, background = [1.0, 1.0, 1.0, 1.0])
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[1])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[1]),
                    self._golden_path.joinpath(TEST_NAMES[1]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[1] + ".diff.png")),
                0.1)

        # render the stereo waveform
        raw = self._audio.draw_waveform(self._prim_2ch, W, H)
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[2])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[2]),
                    self._golden_path.joinpath(TEST_NAMES[2]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[2] + ".diff.png")),
                0.1)

        # add custom channel colors, but only set it for 4 channel to verify that the default still works
        raw = self._audio.draw_waveform(self._prim_2ch, W, H, colors = [[1.0, 0, 1.0, 1.0]])
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[3])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[3]),
                    self._golden_path.joinpath(TEST_NAMES[3]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[3] + ".diff.png")),
                0.1)

        # try 2 custom channel colors now
        raw = self._audio.draw_waveform(self._prim_2ch, W, H, colors = [[1.0, 0, 0, 1.0], [0, 0, 1.0, 1.0]])
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[4])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[4]),
                    self._golden_path.joinpath(TEST_NAMES[4]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[4] + ".diff.png")),
                0.1)

        # this sound asset is 0.693 seconds, so 0.2 is substantial
        self._prim_1ch.GetAttribute("mediaOffsetStart").Set(0.2);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_1ch)

        raw = self._audio.draw_waveform(self._prim_1ch, W, H)
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[5])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[5]),
                    self._golden_path.joinpath(TEST_NAMES[5]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[5] + ".diff.png")),
                0.1)

        self._prim_1ch.GetAttribute("mediaOffsetEnd").Set(0.4);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_1ch)

        raw = self._audio.draw_waveform(self._prim_1ch, W, H)
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[6])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[6]),
                    self._golden_path.joinpath(TEST_NAMES[6]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[6] + ".diff.png")),
                0.1)

        # reset prim parameters
        self._prim_1ch.GetAttribute("mediaOffsetStart").Set(0.0);
        self._prim_1ch.GetAttribute("mediaOffsetEnd").Set(0.0);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_1ch)

        # test it with no flags
        raw = self._audio.draw_waveform(self._prim_1ch, W, H, flags = 0)
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[7])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[7]),
                    self._golden_path.joinpath(TEST_NAMES[7]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[7] + ".diff.png")),
                0.1)

        # test it with alpha blending
        raw = self._audio.draw_waveform(self._prim_1ch, W, H, flags = carb.audio.AUDIO_IMAGE_FLAG_ALPHA_BLEND, colors = [[1.0, 1.0, 1.0, 0.2]])
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[8])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[8]),
                    self._golden_path.joinpath(TEST_NAMES[8]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[8] + ".diff.png")),
                0.1)

        # test choosing the second channel in a multi-channel image
        raw = self._audio.draw_waveform(self._prim_2ch, W, H, flags = 0, channel = 1)
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[9])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[9]),
                    self._golden_path.joinpath(TEST_NAMES[9]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[9] + ".diff.png")),
                0.1)

        # test choosing the second channel in a multi-channel image
        raw = self._audio.draw_waveform(self._prim_2ch, W, H, flags = carb.audio.AUDIO_IMAGE_FLAG_MULTI_CHANNEL | carb.audio.AUDIO_IMAGE_FLAG_ALPHA_BLEND, colors = [[1.0, 0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0]])
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            img.convert("RGB").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[10])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[10]),
                    self._golden_path.joinpath(TEST_NAMES[10]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[10] + ".diff.png")),
                0.1)

        # the colors are going to be random, so turn all the pixels to white before comparing
        raw = self._audio.draw_waveform(self._prim_2ch, W, H, flags = carb.audio.AUDIO_IMAGE_FLAG_NOISE_COLOR)
        self.assertEqual(len(raw), W * H * 4)
        with Image.frombytes("RGBX", (W, H), bytes(raw), 'raw') as img:
            # increase contrast by a factor of 256.0
            ImageEnhance.Contrast(img.convert("RGB")).enhance(256.0).convert("1").save(str(pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[11])))

        if not GENERATE_GOLDEN_IMAGES:
            self.assertLess(omni.kit.test_helpers_gfx.compare_utils.compare(
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[11]),
                    self._golden_path.joinpath(TEST_NAMES[11]),
                    pathlib.Path(OUTPUTS_DIR).joinpath(TEST_NAMES[11] + ".diff.png")),
                2.0)

        # this won't crash or fail, but which flag is chosen is undefined
        raw = self._audio.draw_waveform(self._prim_2ch, W, H, flags = carb.audio.AUDIO_IMAGE_FLAG_MULTI_CHANNEL | carb.audio.AUDIO_IMAGE_FLAG_SPLIT_CHANNELS)
        self.assertEqual(len(raw), W * H * 4)

    async def test_sound_length(self):
        FUZZ = 0.00001
        ASSET_LENGTH = 120.0
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long) - ASSET_LENGTH), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.PLAY_LENGTH) - ASSET_LENGTH), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.SOUND_LENGTH) - ASSET_LENGTH), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.ASSET_LENGTH) - ASSET_LENGTH), FUZZ);

        # start time shouldn't change anything
        self._prim_long.GetAttribute("startTime").Set(12.0);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_long)


        self.assertLess(abs(self._audio.get_sound_length(self._prim_long) - ASSET_LENGTH), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.PLAY_LENGTH) - ASSET_LENGTH), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.SOUND_LENGTH) - ASSET_LENGTH), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.ASSET_LENGTH) - ASSET_LENGTH), FUZZ);

        # end time that exceeds the asset length shouldn't change anything
        self._prim_long.GetAttribute("endTime").Set(ASSET_LENGTH + 20.0);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_long)

        self.assertLess(abs(self._audio.get_sound_length(self._prim_long) - ASSET_LENGTH), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.PLAY_LENGTH) - ASSET_LENGTH), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.SOUND_LENGTH) - ASSET_LENGTH), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.ASSET_LENGTH) - ASSET_LENGTH), FUZZ);

        # end time 6 seconds after start time
        self._prim_long.GetAttribute("endTime").Set(18.0);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_long)

        self.assertLess(abs(self._audio.get_sound_length(self._prim_long) - 6.0), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.PLAY_LENGTH) - 6.0), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.SOUND_LENGTH) - ASSET_LENGTH), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.ASSET_LENGTH) - ASSET_LENGTH), FUZZ);

        # trim the audio track itself
        start_trim = 30.0
        trimmed_length = ASSET_LENGTH - start_trim
        self._prim_long.GetAttribute("mediaOffsetStart").Set(start_trim);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_long)

        self.assertLess(abs(self._audio.get_sound_length(self._prim_long) - 6.0), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.PLAY_LENGTH) - 6.0), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.SOUND_LENGTH) - trimmed_length), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.ASSET_LENGTH) - ASSET_LENGTH), FUZZ);

        # excessive end offset has no effect
        self._prim_long.GetAttribute("mediaOffsetEnd").Set(ASSET_LENGTH + 1.0);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_long)

        self.assertLess(abs(self._audio.get_sound_length(self._prim_long) - 6.0), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.PLAY_LENGTH) - 6.0), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.SOUND_LENGTH) - trimmed_length), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.ASSET_LENGTH) - ASSET_LENGTH), FUZZ);

        end_trim = 32.0
        trimmed_length = end_trim - start_trim

        # use a real trim now
        self._prim_long.GetAttribute("mediaOffsetEnd").Set(end_trim);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_long)

        self.assertLess(abs(self._audio.get_sound_length(self._prim_long) - trimmed_length), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.PLAY_LENGTH) - trimmed_length), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.SOUND_LENGTH) - trimmed_length), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.ASSET_LENGTH) - ASSET_LENGTH), FUZZ);

        # speed it up by a factor of 2
        self._prim_long.GetAttribute("timeScale").Set(2.0);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_long)

        self.assertLess(abs(self._audio.get_sound_length(self._prim_long) - trimmed_length / 2), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.PLAY_LENGTH) - trimmed_length / 2), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.SOUND_LENGTH) - trimmed_length / 2), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.ASSET_LENGTH) - ASSET_LENGTH), FUZZ);

        # try adding a loop count to double the time now
        self._prim_long.GetAttribute("loopCount").Set(1);
        self._prim_long.GetAttribute("timeScale").Set(1.0);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_long)

        self.assertLess(abs(self._audio.get_sound_length(self._prim_long) - 2 * trimmed_length), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.PLAY_LENGTH) - 2 * trimmed_length), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.SOUND_LENGTH) - trimmed_length), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.ASSET_LENGTH) - ASSET_LENGTH), FUZZ);

        # try looping infinitely and remove end time
        self._prim_long.GetAttribute("loopCount").Set(-1);
        self._prim_long.GetAttribute("endTime").Set(-1.0);
        await wait_for_prim_change(self, self._context, self._audio, self._prim_long)

        # check that it's some enormous number
        self.assertGreater(self._audio.get_sound_length(self._prim_long), 2**24);
        self.assertGreater(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.PLAY_LENGTH), 2**24);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.SOUND_LENGTH) - trimmed_length), FUZZ);
        self.assertLess(abs(self._audio.get_sound_length(self._prim_long, omni.usd.audio.SoundLengthType.ASSET_LENGTH) - ASSET_LENGTH), FUZZ);


class TestCaptureStreamer(omni.kit.test.AsyncTestCase):
    # Before running each test
    async def setUp(self):
        extension_path = carb.tokens.get_tokens_interface().resolve("${omni.kit.audio.test.usd}")
        self._usd_path = pathlib.Path(extension_path).joinpath("data").joinpath("tests").absolute()

        from omni.kit.property.usd.usd_attribute_widget import UsdPropertiesWidget
        import omni.kit.window.property as p

        self._w = p.get_window()

        try:
            os.mkdir(OUTPUTS_DIR)
        except FileExistsError:
            pass

    async def test_capture_streamer(self):
        # ****** test setup ******
        audio = omni.usd.audio.get_stage_audio_interface()
        self.assertIsNotNone(audio)
        usd_context = omni.usd.get_context()
        self.assertIsNotNone(usd_context)

        # load a test stage that has some audio prims in it.
        test_file_path = self._usd_path.joinpath("audio_test.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await omni.kit.app.get_app().next_update_async()


        # ****** create and setup the streamers ******
        # create a streamer to capture the audio from the stage.
        streamer = audio.create_capture_streamer()
        self.assertNotEqual(streamer, audio.INVALID_STREAMER_ID)

        streamer2 = audio.create_capture_streamer()
        self.assertNotEqual(streamer2, audio.INVALID_STREAMER_ID)
        self.assertNotEqual(streamer, streamer2)

        streamer3 = audio.create_capture_streamer()
        self.assertNotEqual(streamer3, audio.INVALID_STREAMER_ID)
        self.assertNotEqual(streamer, streamer3)

        serial = random.randrange(2**32)
        def gen_name(name):
            return str(OUTPUTS_DIR.joinpath(str(serial) + "-" + name))

        audio.set_capture_filename(streamer, gen_name("streamer1.wav"))
        audio.set_capture_filename(streamer2, gen_name("streamer2.wav"))
        # intentionally don't set the filename for streamer3.


        # ****** test some captures ******
        # perform a capture on a single streamer with an overridden filename.
        success = audio.start_capture(streamer, gen_name("streamer3.wav"))
        self.assertTrue(success)

        await asyncio.sleep(0.5)
        success = audio.stop_capture(streamer)
        self.assertTrue(success)
        self.assertTrue(audio.wait_for_capture(streamer, 10000))

        self.assertTrue(os.path.exists(gen_name("streamer3.wav")))
        self.assertTrue(os.path.isfile(gen_name("streamer3.wav")))
        self.assertFalse(os.path.exists(gen_name("streamer1.wav")))
        self.assertFalse(os.path.exists(gen_name("streamer2.wav")))

        streamer3_mtime = os.path.getmtime(gen_name("streamer3.wav"))
        await asyncio.sleep(0.5)
        self.assertEqual(streamer3_mtime, os.path.getmtime(gen_name("streamer3.wav")))

        # perform a capture on multiple streamers simultaneously.
        success = audio.start_captures([streamer, streamer2])
        self.assertTrue(success)

        await asyncio.sleep(0.5)
        success = audio.stop_captures([streamer2, streamer])
        self.assertTrue(success)
        self.assertTrue(audio.wait_for_capture(streamer, 10000))
        self.assertTrue(audio.wait_for_capture(streamer2, 10000))

        self.assertTrue(os.path.exists(gen_name("streamer1.wav")))
        self.assertTrue(os.path.exists(gen_name("streamer2.wav")))
        self.assertTrue(os.path.isfile(gen_name("streamer1.wav")))
        self.assertTrue(os.path.isfile(gen_name("streamer2.wav")))

        # streamer3 shouldn't have been modified
        self.assertEqual(streamer3_mtime, os.path.getmtime(gen_name("streamer3.wav")))

        streamer1_mtime = os.path.getmtime(gen_name("streamer1.wav"))
        streamer2_mtime = os.path.getmtime(gen_name("streamer2.wav"))

        # perform a capture on multiple streamers simultaneously where one has not had a filename set.
        audio.set_capture_filename(streamer, gen_name("streamer4.wav"))
        success = audio.start_captures([streamer, streamer3])
        self.assertTrue(success)

        await asyncio.sleep(0.5)
        success = audio.stop_captures([streamer3, streamer])
        self.assertTrue(audio.wait_for_capture(streamer3, 10000))
        self.assertTrue(audio.wait_for_capture(streamer, 10000))

        self.assertTrue(os.path.exists(gen_name("streamer4.wav")))
        self.assertEqual(streamer1_mtime, os.path.getmtime(gen_name("streamer1.wav")))
        self.assertEqual(streamer2_mtime, os.path.getmtime(gen_name("streamer2.wav")))
        self.assertTrue(os.path.isfile(gen_name("streamer4.wav")))
        self.assertEqual(streamer3_mtime, os.path.getmtime(gen_name("streamer3.wav")))


        # start a capture while another streamer is already running.
        audio.set_capture_filename(streamer, gen_name("streamer5.wav"))
        audio.set_capture_filename(streamer2, gen_name("streamer6.wav"))
        success = audio.start_capture(streamer, None)
        self.assertTrue(success)

        await asyncio.sleep(0.5)
        success = audio.start_captures([streamer2])
        self.assertTrue(success)

        await asyncio.sleep(0.5)
        success = audio.stop_captures([streamer2, streamer])
        self.assertTrue(audio.wait_for_capture(streamer, 10000))
        self.assertTrue(audio.wait_for_capture(streamer2, 10000))

        self.assertTrue(os.path.exists(gen_name("streamer5.wav")))
        self.assertTrue(os.path.exists(gen_name("streamer6.wav")))
        self.assertTrue(os.path.isfile(gen_name("streamer5.wav")))
        self.assertTrue(os.path.isfile(gen_name("streamer6.wav")))


        # start a capture on an invalid filename.
        success = audio.start_capture(streamer, "path/does/not/exist/streamer7.wav")
        self.assertFalse(success)

        if success:
            await asyncio.sleep(0.5)
            success = audio.stop_capture(streamer)
            self.assertFalse(success)
            self.assertFalse(os.path.exists("path/does/not/exist/streamer7.wav"))

        success = audio.set_capture_filename(streamer, "other/bad/path/streamer.wav")
        self.assertFalse(success)

        success = audio.set_capture_filename(streamer, "other/bad/path")
        self.assertFalse(success)

        success = audio.set_capture_filename(streamer, "")
        self.assertFalse(success)

        # try starting the streamer before destroying them
        success = audio.start_capture(streamer2, None)
        self.assertTrue(success)

        await asyncio.sleep(0.5)

        # ****** clean up ******
        # destroy the streamers.
        audio.destroy_capture_streamer(streamer)
        audio.destroy_capture_streamer(streamer2)
        audio.destroy_capture_streamer(streamer3)

        # try starting a capture on a destroyed streamer.
        success = audio.start_capture(streamer, gen_name("streamer8.wav"))
        self.assertFalse(success)
        self.assertFalse(os.path.exists(gen_name("streamer8.wav")))

        success = audio.start_capture(streamer2, gen_name("streamer9.wav"))
        self.assertFalse(success)
        self.assertFalse(os.path.exists(gen_name("streamer9.wav")))

        # Need to wait for an additional frames for omni.ui rebuild to take effect
        await wait_stage_loading()

        await usd_context.new_stage_async()
        await omni.kit.app.get_app().next_update_async()

class TestCaptureEvents(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        self._context = omni.usd.get_context()
        self.assertIsNotNone(self._context)

        await self._context.new_stage_async()
        self._stage = self._context.get_stage()
        self.assertIsNotNone(self._stage)

        self._audio = omni.usd.audio.get_stage_audio_interface()
        self.assertIsNotNone(self._audio)

        extension_path = carb.tokens.get_tokens_interface().resolve("${omni.kit.audio.test.usd}")
        self._test_path = pathlib.Path(extension_path).joinpath("data").joinpath("tests").absolute()

        self._test_sound_1ch_path = self._test_path.joinpath("short-1ch.oga").absolute()
        self._prim_1ch_path = "/sound_1ch"
        self._prim_1ch = self._stage.DefinePrim(self._prim_1ch_path, "Sound")
        self.assertIsNotNone(self._prim_1ch)
        self._prim_1ch.GetAttribute("filePath").Set(str(self._test_sound_1ch_path))
        self._prim_1ch.GetAttribute("startTime").Set(-1.0)

        self._streamer = self._audio.create_capture_streamer()
        self.assertNotEqual(self._streamer, omni.usd.audio.INVALID_STREAMER_ID)

    async def tearDown(self):
        self._audio.destroy_capture_streamer(self._streamer)

    async def test_capture_events(self):

        fmt = None
        open_count = 0
        close_count = 0
        data = []
        def open_impl(stream_fmt):
            nonlocal fmt
            nonlocal open_count
            fmt = stream_fmt
            open_count += 1

        def write_impl(stream_data):
            nonlocal data
            data += stream_data

        def close_impl():
            nonlocal close_count
            close_count += 1

        events = self._audio.create_event_stream_for_capture(self._streamer)
        self.assertIsNotNone(events)
        self.assertTrue(events)

        # create the listener in the middle of the stream and verify that nothing
        # is actually produced
        self._audio.start_capture(self._streamer)
        await asyncio.sleep(0.2)

        listener = omni.usd.audio.StreamListener(events, open_impl, write_impl, close_impl)
        self.assertIsNotNone(listener)

        self._audio.stop_capture(self._streamer)
        await asyncio.sleep(0.2)

        self.assertIsNone(fmt)
        self.assertEqual(open_count, 0);
        self.assertEqual(close_count, 0);
        self.assertEqual(len(data), 0);

        # this time, it should actually do something
        self._audio.start_capture(self._streamer)
        await asyncio.sleep(0.2)

        self.assertEqual(open_count, 1);
        self.assertIsNotNone(fmt)
        self.assertEqual(close_count, 0);
        self.assertGreater(len(data), 0);

        self._audio.stop_capture(self._streamer)
        await asyncio.sleep(0.2)

        self.assertEqual(open_count, 1);
        self.assertIsNotNone(fmt)
        self.assertEqual(close_count, 1);
        self.assertGreater(len(data), 0);

        # check that nothing was played
        silence = 0
        if fmt.format == carb.audio.SampleFormat.PCM8:
            silence = 127

        for v in data:
            self.assertEqual(v, silence);


        # this time we'll actually play a sound with it
        data = []
        self._audio.start_capture(self._streamer)
        await asyncio.sleep(0.1)
        voice = self._audio.spawn_voice(self._prim_1ch)
        i = 0
        while voice.is_playing():
            await asyncio.sleep(0.01)
            i += 1
            if (i == 2000):
                self.assertFalse("test timed out")

        self._audio.stop_capture(self._streamer)
        await asyncio.sleep(0.2)

        self.assertEqual(open_count, 2);
        self.assertIsNotNone(fmt)
        self.assertEqual(close_count, 2);
        self.assertGreater(len(data), 0);

        silent = True
        for v in data:
            if v != silence:
                silent = False
                break

        self.assertFalse(silent)


class TestListenerEnum(omni.kit.test.AsyncTestCase):
    # Before running each test
    async def setUp(self):
        extension_path = carb.tokens.get_tokens_interface().resolve("${omni.kit.audio.test.usd}")
        self._usd_path = pathlib.Path(extension_path).joinpath("data").joinpath("tests").absolute()

        from omni.kit.property.usd.usd_attribute_widget import UsdPropertiesWidget
        import omni.kit.window.property as p

        self._w = p.get_window()

    async def test_listener_enum(self):
        # ****** test setup ******
        audio = omni.usd.audio.get_stage_audio_interface()
        self.assertIsNotNone(audio)
        usd_context = omni.usd.get_context()
        self.assertIsNotNone(usd_context)

        # load a test stage that has some audio prims in it.
        test_file_path = self._usd_path.joinpath("audio_test2.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await omni.kit.app.get_app().next_update_async()
        await wait_stage_loading()


        # ****** try some bad parameters ******
        count = audio.get_listener_count()
        self.assertNotEqual(count, 0)

        # try an out of range index.  Note that we can't try a negative number here because
        # pybind freaks out and throws a TypeError exception because it doesn't match the
        # `size_t` argument.
        prim = audio.get_listener_by_index(count)
        self.assertIsNone(prim)

        prim = audio.get_listener_by_index(count + 1)
        self.assertIsNone(prim)

        prim = audio.get_listener_by_index(count + 20000)
        self.assertIsNone(prim)

        prim = audio.get_listener_by_index(count - 1)
        self.assertIsNotNone(prim)

        prim = None


        # ****** enumerate the stage listeners ******
        count = audio.get_listener_count()
        self.assertNotEqual(count, 0)

        for i in range(0, count):
            prim = audio.get_listener_by_index(i)
            self.assertIsNotNone(prim)


        # Need to wait for an additional frames for omni.ui rebuild to take effect
        await wait_stage_loading()

        await usd_context.new_stage_async()
        await omni.kit.app.get_app().next_update_async()

class TestStageAudioListeners(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        self._context = omni.usd.get_context()
        self.assertIsNotNone(self._context)

        self._context.new_stage()
        self._stage = self._context.get_stage()
        self.assertIsNotNone(self._stage)

        self._audio = omni.usd.audio.get_stage_audio_interface()
        self.assertIsNotNone(self._audio)

        extension_path = carb.tokens.get_tokens_interface().resolve("${omni.kit.audio.test.usd}")

        self._prim_path = "/listener"
        self._prim = self._stage.DefinePrim(self._prim_path, "Listener")
        self.assertIsNotNone(self._prim)

        await self._context.next_frame_async()

    async def tearDown(self):
        self._stage.RemovePrim(self._prim_path)

    async def test_get_listener_count(self):
        self.assertEqual(self._audio.get_listener_count(), 1)

        self._new_prim_path = "/listener1"
        self._new_prim = self._stage.DefinePrim(self._new_prim_path, "Listener")

        await self._context.next_frame_async()
        self.assertEqual(self._audio.get_listener_count(), 2)

        self._stage.RemovePrim(self._new_prim_path)

    async def test_set_active_listener(self):
        self.assertIsNone(self._audio.get_active_listener())

        self._new_prim_path = "/sound"
        self._new_prim = self._stage.DefinePrim(self._new_prim_path, "Sound")
        await self._context.next_frame_async()

        self.assertFalse(self._audio.set_active_listener(self._new_prim))

        self._stage.RemovePrim(self._new_prim_path)

        self.assertTrue(self._audio.set_active_listener(self._prim))
        self.assertEqual(self._audio.get_active_listener(), self._prim)

        self.assertTrue(self._audio.set_active_listener(None))
        self.assertIsNone(self._audio.get_active_listener())

        # FIXME: This is just verifying that we can get/set this value.
        #        This test should somehow verify that this feature is actually
        #        detectable in the output audio.

    async def test_get_listener_by_index(self):
        # full enumeration of the 1 sound in the scene
        self.assertEqual(self._audio.get_listener_by_index(0), self._prim)
        self.assertIsNone(self._audio.get_listener_by_index(1))

