import pytest
from spharpy.classes import SphericalHarmonicSignal
from spharpy.classes import SphericalHarmonicDefinition
import numpy as np
import re


def test_spherical_harmonic_signal_init():
    """Test init SphericalHarmonicsSignal."""

    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 4, 3)
    signal = SphericalHarmonicSignal(data,
                                     44100, basis_type='real',
                                     channel_convention='ACN',
                                     normalization='N3D',
                                     condon_shortley=False)
    assert isinstance(signal, SphericalHarmonicSignal)


def test_sh_signal_from_sh_definition():
    """Test init SphericalHarmonicsSignal."""

    shd = SphericalHarmonicDefinition(
        n_max=0, basis_type="real",
        channel_convention="ACN",
        normalization="N3D",
        condon_shortley=False)

    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 4, 3)

    signal_def = SphericalHarmonicSignal.from_definition(
        sh_definition=shd, data=data, sampling_rate=44100)

    assert isinstance(signal_def, SphericalHarmonicSignal)

    signal = SphericalHarmonicSignal(data,
                                     44100, basis_type='real',
                                     channel_convention='ACN',
                                     normalization='N3D',
                                     condon_shortley=False)
    assert signal_def == signal


def test_spherical_harmonic_signal_init_condon_shortley():
    """
    Test if Condon-Shortley is set properly in init
    SphericalHarmonicsSignal.
    """

    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 4, 3)
    signal = SphericalHarmonicSignal(data,
                                     44100, basis_type='real',
                                     channel_convention='ACN',
                                     condon_shortley=False,
                                     normalization='N3D')
    assert not signal.condon_shortley

    signal = SphericalHarmonicSignal(data.astype(np.complex128),
                                     44100, basis_type='complex',
                                     channel_convention='ACN',
                                     normalization='N3D',
                                     condon_shortley=True,
                                     is_complex=True)
    assert signal.condon_shortley


def test_spherical_harmonic_signal_wrong_dimensions():
    """Test dimensions of SH coefficient data."""

    # test if sh channels are valid
    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 5, 3)

    with pytest.raises(ValueError,
                   match=re.escape("Invalid number of spherical harmonic "
                                   f"channels: {data.shape[-2]}. It must "
                                   "match (n_max + 1)^2.")):
        SphericalHarmonicSignal(data,
                                44100, basis_type='real',
                                channel_convention='ACN',
                                condon_shortley=False,
                                normalization='N3D')


def test_spherical_harmonic_signal_init_multichannel():
    """Test init SphercalHarmonicsSignal."""
    sh_coeffs = np.zeros((2, 4, 16))
    signal = SphericalHarmonicSignal(sh_coeffs,
                                     44100, basis_type='real',
                                     channel_convention='ACN',
                                     normalization='N3D',
                                     condon_shortley=False)
    assert isinstance(signal, SphericalHarmonicSignal)


def test_nmax_getter():
    """Test nmax getter."""
    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 4, 3)
    signal = SphericalHarmonicSignal(data,
                                     44100, basis_type='real',
                                     channel_convention='ACN',
                                     normalization='N3D',
                                     condon_shortley=False)
    assert signal.n_max == 1
    assert isinstance(signal.n_max, int)


def test_init_wrong_basis_type():
    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 4, 3)
    with pytest.raises(ValueError,
                match="Invalid basis type, only "
                      "'complex' and 'real' are supported"):
        SphericalHarmonicSignal(data,
                                44100, basis_type='invalid_basis_type',
                                channel_convention='ACN',
                                normalization='N3D',
                                condon_shortley=False)


def test_basis_type_getter():
    """Test basis_type getter."""
    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 4, 3)
    signal = SphericalHarmonicSignal(data,
                                     44100, basis_type='real',
                                     channel_convention='ACN',
                                     normalization='N3D',
                                     condon_shortley=False)
    assert signal.basis_type == 'real'


def test_init_wrong_normalization():
    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 4, 3)

    with pytest.raises(ValueError,
                       match="Invalid normalization, currently only 'N3D', "
                             "'NM', 'maxN', 'SN3D', 'SNM' are supported"):
        SphericalHarmonicSignal(data,
                                44100, basis_type='real',
                                channel_convention='ACN',
                                normalization='invalid_normalization',
                                condon_shortley=False)


def test_spherical_harmonic_signal_normalization_setter():
    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 4, 3)
    signal = SphericalHarmonicSignal(data,
                                     44100, basis_type='real',
                                     channel_convention='ACN',
                                     normalization='N3D',
                                     condon_shortley=False)
    signal.normalization = 'SN3D'
    assert signal.normalization == 'SN3D'

    signal.normalization = 'SNM'
    assert signal.normalization == 'SNM'

    signal.normalization = 'maxN'
    assert signal.normalization == 'maxN'

    signal.normalization = 'N3D'
    assert signal.normalization == 'N3D'

    signal.normalization = 'NM'
    assert signal.normalization == 'NM'


def test_init_wrong_channel_convention():
    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 4, 3)

    with pytest.raises(
        ValueError,
        match="Invalid channel convention, currently only 'ACN' "
              "and 'FuMa' are supported"):
        SphericalHarmonicSignal(data,
                                44100, basis_type='real',
                                channel_convention='invalid_convention',
                                normalization='N3D',
                                condon_shortley=False)


def test_spherical_harmonic_signal_channel_convention_setter():
    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 4, 3)
    signal = SphericalHarmonicSignal(data,
                                     44100, basis_type='real',
                                     channel_convention='ACN',
                                     normalization='N3D',
                                     condon_shortley=False)
    signal.channel_convention = 'FuMa'
    assert signal.channel_convention == 'FuMa'

    signal.channel_convention = 'ACN'
    assert signal.channel_convention == 'ACN'


def test_spherical_harmonic_signal_change_channel_convention():
    data = np.array([[1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.],
                     [1., 2., 3.]]).reshape(1, 4, 3)
    signal = SphericalHarmonicSignal(data,
                                     44100, basis_type='real',
                                     channel_convention='ACN',
                                     normalization='N3D',
                                     condon_shortley=False)
    signal.channel_convention = 'FuMa'
    assert signal.channel_convention == 'FuMa'

    data_fuma = data[:, [0, 3, 1, 2], :]

    np.testing.assert_equal(data_fuma,
                            signal.time)


def test_spherical_harmonic_signal_renormalize():
    data = np.ones((1, 4, 2))
    signal = SphericalHarmonicSignal(data,
                                     44100, basis_type='real',
                                     channel_convention='ACN',
                                     normalization='N3D',
                                     condon_shortley=False)
    signal.normalization = 'maxN'
    data_ref = np.array([[np.sqrt(1 / 2), np.sqrt(1 / 2)],
                         [np.sqrt(1 / 3), np.sqrt(1 / 3)],
                         [np.sqrt(1 / 3), np.sqrt(1 / 3)],
                         [np.sqrt(1 / 3), np.sqrt(1 / 3)]]).reshape(1, 4, 2)

    np.testing.assert_equal(signal.time,
                            data_ref)
