#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Tests for edc related things."""

import numpy as np
from numpy import array

import numpy.testing as npt
import pyfar as pf
import pytest

import pyrato as ra


@pytest.mark.parametrize("is_energy", [True, False])
def test_schroeder_integration(is_energy):
    """Simple impulse response."""
    factor = 2 if is_energy else np.sqrt(2)
    data = pf.Signal(np.ones((1, 10))*factor, sampling_rate=10)
    energy_decay_curve = ra.edc.schroeder_integration(
        data, is_energy=is_energy)

    # test if the shape is correct
    npt.assert_array_equal(energy_decay_curve.time.shape, data.time.shape)

    # test if the values are correct
    truth = (np.arange(10, dtype=float)[::-1] + 1) * 2
    npt.assert_almost_equal(energy_decay_curve.time[0], truth)


@pytest.mark.parametrize("is_energy", [True, False])
def test_schroeder_integration_multi_cdim(is_energy):
    """Simple impulse response."""
    factor = 2 if is_energy else np.sqrt(2)
    data = pf.Signal(np.ones((2, 1, 10))*factor, sampling_rate=10)
    energy_decay_curve = ra.edc.schroeder_integration(
        data, is_energy=is_energy)

    # test if the shape is correct
    npt.assert_array_equal(energy_decay_curve.time.shape, data.time.shape)

    # test if the values are correct
    truth = (np.arange(10, dtype=float)[::-1] + 1) * 2
    npt.assert_almost_equal(energy_decay_curve.time[0, 0], truth)
    npt.assert_almost_equal(energy_decay_curve.time[1, 0], truth)


@pytest.mark.parametrize(
        "edc_function",
        [ra.edc.energy_decay_curve_chu,
         ra.edc.energy_decay_curve_lundeby,
         ra.edc.energy_decay_curve_chu_lundeby],
)
def test_multidim_edc(edc_function):
    """
    Test if edcs from multichannel signal are equal to corresponding single
    channel edcs.
    """
    rir = pf.signals.files.room_impulse_response()
    rir_oct = pf.dsp.filter.fractional_octave_bands(rir, 1)
    shape = rir_oct.time.shape
    edc = edc_function(rir_oct, channel_independent=True)

    assert shape == edc.time.shape

    edc = edc.flatten()
    rir_oct = rir_oct.flatten()

    for i in range(edc.cshape[0]):
        baseline = edc_function(rir_oct[i])
        npt.assert_array_equal(edc[i].time, baseline.time)


def test_edc_function():
    """Test for the parametric energy decay curve based on previous
    implementation which also included RT prediction using Sabine's equation.
    """
    times = np.linspace(0, 0.25, 50)

    c = 343.4

    volume = 2*2*2
    alphas = np.asarray([0.9, 0.1])
    surfaces = np.asarray([2, 5*2])
    surface_room = np.sum(surfaces)
    alpha_mean = np.sum(surfaces*alphas) / surface_room

    rt = 24*np.log(10)/c * volume / (surface_room*alpha_mean)
    energy = 1
    edc = ra.parametric.energy_decay_curve(times, rt, energy)

    truth = array([
        1.00000000e+00, 8.57869258e-01, 7.35939663e-01, 6.31340013e-01,
        5.41607188e-01, 4.64628156e-01, 3.98590212e-01, 3.41938289e-01,
        2.93338346e-01, 2.51645949e-01, 2.15879324e-01, 1.85196235e-01,
        1.58874157e-01, 1.36293255e-01, 1.16921793e-01, 1.00303612e-01,
        8.60473853e-02, 7.38174065e-02, 6.33256837e-02, 5.43251573e-02,
        4.66038824e-02, 3.99800380e-02, 3.42976455e-02, 2.94228957e-02,
        2.52409977e-02, 2.16534759e-02, 1.85758513e-02, 1.59356518e-02,
        1.36707058e-02, 1.17276782e-02, 1.00608146e-02, 8.63086356e-03,
        7.40415251e-03, 6.35179482e-03, 5.44900951e-03, 4.67453774e-03,
        4.01014222e-03, 3.44017773e-03, 2.95122272e-03, 2.53176324e-03,
        2.17192185e-03, 1.86322499e-03, 1.59840344e-03, 1.37122117e-03,
        1.17632849e-03, 1.00913605e-03, 8.65706791e-04, 7.42663242e-04,
        6.37107964e-04, 5.46555336e-04,
        ])
    npt.assert_almost_equal(edc.time, np.atleast_2d(truth))


def test_parametric_edc():
    """Test returned edc by the parametric edc function and shape broadcasting.
    """
    times = np.linspace(0, 0.25, 50)
    T_60 = np.array([2, 1])

    edc = ra.parametric.energy_decay_curve(times, T_60)
    assert edc.cshape == (2,)

    e_0 = np.ones_like(T_60)
    edc = ra.parametric.energy_decay_curve(times, T_60, energy=e_0)

    damping_term = 6*np.log(10) / T_60
    truth = np.atleast_2d(e_0).T * np.exp(-np.atleast_2d(damping_term).T*times)

    npt.assert_almost_equal(edc.time, truth)
    assert edc.cshape == (2,)

    edc = ra.parametric.energy_decay_curve(
        times, np.ones((3, 2)), energy=np.ones((3, 2)))

    assert edc.cshape == (3, 2)


def test_parametric_edc_wrong_shapes():
    """Test error handling for wrong shapes."""
    times = np.linspace(0, 0.25, 50)
    T_60 = np.array([2, 1])
    energy = np.array([1, 1, 1])

    with pytest.raises(ValueError, match="same shape."):
        ra.parametric.energy_decay_curve(times, T_60, energy=energy)
