import numpy as np
from numpy.testing import assert_almost_equal

from audiomentations import Shift, Compose


class TestShift:
    def test_shift(self):
        samples = np.array([1.0, 0.5, 0.25, 0.125], dtype=np.float32)
        sample_rate = 16000

        forward_augmenter = Compose([Shift(min_fraction=0.5, max_fraction=0.5, p=1.0)])
        forward_shifted_samples = forward_augmenter(
            samples=samples, sample_rate=sample_rate
        )
        assert_almost_equal(
            forward_shifted_samples, np.array([0.25, 0.125, 1.0, 0.5], dtype=np.float32)
        )
        assert forward_shifted_samples.dtype == np.float32
        assert len(forward_shifted_samples) == 4

        backward_augmenter = Compose(
            [Shift(min_fraction=-0.25, max_fraction=-0.25, p=1.0)]
        )
        backward_shifted_samples = backward_augmenter(
            samples=samples, sample_rate=sample_rate
        )
        assert_almost_equal(
            backward_shifted_samples,
            np.array([0.5, 0.25, 0.125, 1.0], dtype=np.float32),
        )
        assert backward_shifted_samples.dtype == np.float32
        assert len(forward_shifted_samples) == 4

    def test_shift_without_rollover(self):
        samples = np.array([1.0, 0.5, 0.25, 0.125], dtype=np.float32)
        sample_rate = 16000

        forward_augmenter = Compose(
            [Shift(min_fraction=0.5, max_fraction=0.5, rollover=False, p=1.0)]
        )
        forward_shifted_samples = forward_augmenter(
            samples=samples, sample_rate=sample_rate
        )
        assert_almost_equal(
            forward_shifted_samples, np.array([0.0, 0.0, 1.0, 0.5], dtype=np.float32)
        )
        assert forward_shifted_samples.dtype == np.float32
        assert len(forward_shifted_samples) == 4

        backward_augmenter = Compose(
            [Shift(min_fraction=-0.25, max_fraction=-0.25, rollover=False, p=1.0)]
        )
        backward_shifted_samples = backward_augmenter(
            samples=samples, sample_rate=sample_rate
        )
        assert_almost_equal(
            backward_shifted_samples,
            np.array([0.5, 0.25, 0.125, 0.0], dtype=np.float32),
        )
        assert backward_shifted_samples.dtype == np.float32
        assert len(forward_shifted_samples) == 4

    def test_shift_multichannel(self):
        samples = np.array(
            [[0.75, 0.5, -0.25, -0.125], [0.9, 0.5, -0.25, -0.125]], dtype=np.float32
        )
        sample_rate = 4000

        augment = Shift(min_fraction=0.5, max_fraction=0.5, p=1.0)
        processed_samples = augment(samples=samples, sample_rate=sample_rate)

        assert_almost_equal(
            processed_samples,
            np.array(
                [[-0.25, -0.125, 0.75, 0.5], [-0.25, -0.125, 0.9, 0.5]],
                dtype=np.float32,
            ),
        )
        assert processed_samples.dtype == np.float32

    def test_shift_without_rollover_multichannel(self):
        samples = np.array(
            [[0.75, 0.5, -0.25, -0.125], [0.9, 0.5, -0.25, -0.125]], dtype=np.float32
        )
        sample_rate = 4000

        augment = Shift(min_fraction=0.5, max_fraction=0.5, rollover=False, p=1.0)
        processed_samples = augment(samples=samples, sample_rate=sample_rate)

        assert_almost_equal(
            processed_samples,
            np.array([[0.0, 0.0, 0.75, 0.5], [0.0, 0.0, 0.9, 0.5]], dtype=np.float32),
        )
        assert processed_samples.dtype == np.float32

    def test_shift_fade(self):
        samples = np.array(
            [[1.0, 2.0, 3.0, 4.0, 5.0], [-1.0, -2.0, -3.0, -4.0, -5.0]],
            dtype=np.float32,
        )
        sample_rate = 4000

        augment = Shift(
            min_fraction=0.5,
            max_fraction=0.5,
            rollover=False,
            fade=True,
            fade_duration=0.00075,  # 0.00075 * 4000 = 3
            p=1.0,
        )
        processed_samples = augment(samples=samples, sample_rate=sample_rate)

        assert_almost_equal(
            processed_samples,
            np.array(
                [[0.0, 0.0, 0.0, 1.0, 3.0], [0.0, 0.0, 0.0, -1.0, -3.0]],
                dtype=np.float32,
            ),
        )

    def test_shift_fade_rollover(self):
        samples = np.array(
            [[1.0, 2.0, 3.0, 4.0, 5.0], [-1.0, -2.0, -3.0, -4.0, -5.0]],
            dtype=np.float32,
        )
        sample_rate = 4000

        augment = Shift(
            min_fraction=0.5,
            max_fraction=0.5,
            rollover=True,
            fade=True,
            fade_duration=0.00075,  # 0.00075 * 4000 = 3
            p=1.0,
        )
        processed_samples = augment(samples=samples, sample_rate=sample_rate)
        assert_almost_equal(
            processed_samples,
            np.array(
                [[2.0, 0.0, 0, 1.0, 3.0], [-2.0, 0.0, 0, -1.0, -3.0]], dtype=np.float32,
            ),
        )

    def test_shift_fade_rollover_2(self):
        samples = np.array(
            [[1.0, 2.0, 3.0, 4.0, 5.0], [-1.0, -2.0, -3.0, -4.0, -5.0]],
            dtype=np.float32,
        )
        sample_rate = 4000

        augment = Shift(
            min_fraction=-0.5,
            max_fraction=-0.5,
            rollover=True,
            fade=True,
            fade_duration=0.00075,  # 0.00075 * 4000 = 3
            p=1.0,
        )
        processed_samples = augment(samples=samples, sample_rate=sample_rate)
        assert_almost_equal(
            processed_samples,
            np.array(
                [[3.0, 2.0, 0.0, 0.0, 1.0], [-3.0, -2.0, 0.0, -0.0, -1.0]],
                dtype=np.float32,
            ),
        )

    def test_shift_fade_rollover_3(self):
        samples = np.array(
            [[1.0, 2.0, 3.0, 4.0, 5.0], [-1.0, -2.0, -3.0, -4.0, -5.0]],
            dtype=np.float32,
        )
        sample_rate = 4000

        augment = Shift(
            min_fraction=-0.5,
            max_fraction=-0.5,
            rollover=True,
            fade=True,
            fade_duration=1.0,
            p=1.0,
        )
        processed_samples = augment(samples=samples, sample_rate=sample_rate)
        assert_almost_equal(
            processed_samples,
            np.array(
                [
                    [0.0015004, 0.0010003, 0.0, 0.0, 0.0005001],
                    [-0.0015004, -0.0010003, -0.0, -0.0, -0.0005001],
                ],
                dtype=np.float32,
            ),
        )

    def test_freeze_shift_with_different_sample_rates(self):
        samples1 = np.array([1.0, 0.5, 0.25, 0.125], dtype=np.float32)
        samples2 = np.array(
            [1.0, 0.5, 0.25, 0.125, 0.3, 0.4, 0.5, 0.2], dtype=np.float32
        )
        sample_rate_1 = 1
        sample_rate_2 = 2

        augment = Shift(
            min_fraction=0.5, max_fraction=0.5, rollover=False, fade=False, p=1.0
        )
        augment.randomize_parameters(samples1, sample_rate_1)
        augment.freeze_parameters()
        shifted_samples1 = augment(samples1, sample_rate_1)
        shifted_samples2 = augment(samples2, sample_rate_2)

        num_samples1_shifted = np.count_nonzero(shifted_samples1 == 0)
        num_samples2_shifted = np.count_nonzero(shifted_samples2 == 0)

        assert num_samples1_shifted == 2
        assert num_samples2_shifted == 4

