import numpy as np
import pytest

from ssspy.algorithm.permutation_alignment import (
    correlation_based_permutation_solver,
    score_based_permutation_solver,
)

rng = np.random.default_rng(0)

parameters_give_demixing_filter = [True, False]


@pytest.mark.parametrize("give_demixing_filter", parameters_give_demixing_filter)
def test_correlation_based_permutation_solver(give_demixing_filter: bool):
    n_sources = 3
    n_channels = n_sources
    n_bins, n_frames = 4, 16

    shape = (n_channels, n_bins, n_frames)
    mixture = rng.standard_normal(shape) + 1j * rng.standard_normal(shape)
    shape = (n_bins, n_sources, n_channels)
    demix_filter = rng.standard_normal(shape) + 1j * rng.standard_normal(shape)
    separated = demix_filter @ mixture.transpose(1, 0, 2)

    if give_demixing_filter:
        separated, demix_filter = correlation_based_permutation_solver(separated, demix_filter)

        assert demix_filter.shape == (n_bins, n_sources, n_channels)
    else:
        separated = correlation_based_permutation_solver(separated)

    assert separated.shape == (n_bins, n_sources, n_frames)


@pytest.mark.parametrize("give_demixing_filter", parameters_give_demixing_filter)
def test_score_based_permutation_solver(give_demixing_filter: bool):
    n_sources = 3
    n_channels = n_sources
    n_bins, n_frames = 4, 16

    shape = (n_channels, n_bins, n_frames)
    mixture = rng.standard_normal(shape) + 1j * rng.standard_normal(shape)
    shape = (n_bins, n_sources, n_channels)
    demix_filter = rng.standard_normal(shape) + 1j * rng.standard_normal(shape)
    separated = demix_filter @ mixture.transpose(1, 0, 2)

    if give_demixing_filter:
        separated, demix_filter = score_based_permutation_solver(separated, demix_filter)

        assert demix_filter.shape == (n_bins, n_sources, n_channels)
    else:
        separated = correlation_based_permutation_solver(separated)

    assert separated.shape == (n_bins, n_sources, n_frames)
