#!/usr/bin/python3
import numpy as np
from lasp import cppSLM
from lasp.filter import SPLFilterDesigner


def test_cppslm1():
    """
    Generate a sine wave
    """
    fs = 48000
    omg = 2 * np.pi * 1000

    slm = cppSLM.fromBiquads(fs, 2e-5, 1, 0.125,
                             np.array([[1., 0, 0, 1, 0, 0]]).T)

    t = np.linspace(0, 10, 10 * fs, endpoint=False)

    # Input signal with an rms of 1 Pa
    in_ = np.sin(omg * t) * np.sqrt(2)

    # Compute overall RMS
    rms = np.sqrt(np.sum(in_**2) / in_.size)
    # Compute overall level
    level = 20 * np.log10(rms / 2e-5)

    # Output of SLM
    out = slm.run(in_)

    # Output of SLM should be close to theoretical
    # level, at least for reasonable time constants
    # (Fast, Slow etc)
    assert (np.isclose(out[-1, 0], level))


def test_cppslm2():
    """
    Generate a sine wave, now A-weighted
    """
    fs = 48000
    omg = 2 * np.pi * 1000

    filt = SPLFilterDesigner(fs).A_Sos_design()
    slm = cppSLM.fromBiquads(fs, 2e-5, 0, 0.125,
                             filt.flatten(),  # Pre-filter coefs
                             np.array([[1., 0, 0, 1, 0, 0]]).T # Bandpass coefs
                            )

    t = np.linspace(0, 10, 10 * fs, endpoint=False)

    # Input signal with an rms of 1 Pa
    in_ = np.sin(omg * t) * np.sqrt(2)

    # Compute overall RMS
    rms = np.sqrt(np.sum(in_**2) / in_.size)
    # Compute overall level
    level = 20 * np.log10(rms / 2e-5)

    # Output of SLM
    out = slm.run(in_)

    # Output of SLM should be close to theoretical
    # level, at least for reasonable time constants
    # (Fast, Slow etc)
    assert np.isclose(out[-1, 0], level, atol=1e-2)


def test_cppslm3():
    fs = 48000
    omg = 2 * np.pi * 1000

    filt = SPLFilterDesigner(fs).A_Sos_design()
    slm = cppSLM.fromBiquads(fs, 2e-5, 0, 0.125,
                             filt.flatten(), 
                             np.array([[1., 0, 0, 1, 0, 0]]).T)
    t = np.linspace(0, 10, 10 * fs, endpoint=False)

    in_ = 10 * np.sin(omg * t) * np.sqrt(2) + np.random.randn()
    # Compute overall RMS
    rms = np.sqrt(np.sum(in_**2) / in_.size)
    # Compute overall level
    level = 20 * np.log10(rms / 2e-5)

    # Output of SLM
    out = slm.run(in_)

    Lpeak = 20 * np.log10(np.max(np.abs(in_) / 2e-5))
    Lpeak
    slm.Lpeak()
    assert np.isclose(out[-1, 0], slm.Leq()[0], atol=1e-2)
    assert np.isclose(Lpeak, slm.Lpeak()[0], atol=2e0)


if __name__ == '__main__':
    test_cppslm1()
    test_cppslm2()
    test_cppslm3()