import soundfile as sf
import jaxloudnorm as jln

data, rate = sf.read("test.wav") # load audio (with shape (samples, channels))
meter = jln.Meter(rate) # create BS.1770 meter
loudness = meter.integrated_loudness(data) # measure loudness

import soundfile as sf
import jaxloudnorm as jln

data, rate = sf.read("test.wav") # load audio

# peak normalize audio to -1 dB
peak_normalized_audio = jln.normalize.peak(data, -1.0)

# measure the loudness first 
meter = jln.Meter(rate) # create BS.1770 meter
loudness = meter.integrated_loudness(data)

# loudness normalize audio to -12 dB LUFS
loudness_normalized_audio = jln.normalize.loudness(data, loudness, -12.0)

import soundfile as sf
import jaxloudnorm as jln
from jaxloudnorm import IIRfilter

data, rate = sf.read("test.wav") # load audio

# block size
meter1 = jln.Meter(rate)                               # 400ms block size
meter2 = jln.Meter(rate, block_size=0.200)             # 200ms block size

# filter classes
meter3 = jln.Meter(rate)                               # BS.1770 meter
meter4 = jln.Meter(rate, filter_class="DeMan")         # fully compliant filters  
meter5 = jln.Meter(rate, filter_class="Fenton/Lee 1")  # low complexity improvement by Fenton and Lee
meter6 = jln.Meter(rate, filter_class="Fenton/Lee 2")  # higher complexity improvement by Fenton and Lee
meter7 = jln.Meter(rate, filter_class="Dash et al.")   # early modification option

# create your own IIR filters
my_high_pass  = IIRfilter(0.0, 0.5, 20.0, rate, 'high_pass')
my_high_shelf = IIRfilter(2.0, 0.7, 1525.0, rate, 'high_shelf')

# create a meter initialized without filters
meter8 = jln.Meter(rate, filter_class="custom")

# load your filters into the meter
meter8._filters = [my_high_pass, my_high_shelf]

def test_batched_integrated_loudness():
    data, rate = sf.read("tests/data/sine_1000.wav")
    meter = pyln.Meter(rate)
    loudness = jax.vmap(meter.integrated_loudness)(jnp.stack([data, data, data]))

    assert jnp.allclose(loudness, jnp.full(loudness.shape, -3.0523438444331137))

def test_batched_loudness_normalize():
    data, rate = sf.read("tests/data/sine_1000.wav")
    data = jnp.stack([data, data, data, data])
    meter = pyln.Meter(rate)
    loudness = jax.vmap(meter.integrated_loudness)(data)
    norm = jax.vmap(pyln.normalize.loudness, in_axes=(0, 0, None))(data, loudness, -6.0)
    loudness = jax.vmap(meter.integrated_loudness)(norm)

    assert jnp.allclose(loudness, jnp.full(loudness.shape, -6.0))