import numpy as np
import pytest

from fasr.data.audio import (
    Audio,
    AudioChannel,
    AudioChannelList,
    AudioSpan,
    AudioSpanList,
    AudioToken,
    AudioTokenList,
)
from fasr.data.eval import AudioEvalExample, AudioEvalResult
from fasr.data.waveform import Waveform


def _make_audio_with_text(text: str, sample_rate: int = 16000) -> Audio:
    data = np.zeros(sample_rate, dtype=np.float32)
    wf = Waveform(data=data, sample_rate=sample_rate)
    ch = AudioChannel(waveform=wf, raw_text=text)
    audio = Audio(sample_rate=sample_rate, duration=1.0)
    audio.channels = AudioChannelList([ch])
    return audio


class TestAudioEvalExample:
    def test_create(self):
        x = _make_audio_with_text("你好世界")
        y = _make_audio_with_text("你好世界")
        example = AudioEvalExample(x=x, y=y)
        assert example.is_badcase is False

    def test_channel_mismatch_raises(self):
        x = _make_audio_with_text("hello")
        y = Audio(sample_rate=16000)
        y.channels = AudioChannelList([AudioChannel(), AudioChannel()])
        with pytest.raises(ValueError):
            AudioEvalExample(x=x, y=y)

    def test_char_error_rate_perfect(self):
        x = _make_audio_with_text("你好世界")
        y = _make_audio_with_text("你好世界")
        example = AudioEvalExample(x=x, y=y)
        cer = example.char_error_rate()
        assert cer == 0.0

    def test_char_error_rate_different(self):
        x = _make_audio_with_text("你好啊世界")
        y = _make_audio_with_text("你好世界")
        example = AudioEvalExample(x=x, y=y)
        cer = example.char_error_rate()
        assert cer > 0
        assert example.is_badcase is True

    def test_clear_waveform(self):
        x = _make_audio_with_text("hello")
        y = _make_audio_with_text("hello")
        example = AudioEvalExample(x=x, y=y)
        example.clear_waveform()
        assert x.channels[0].waveform is None
        assert y.channels[0].waveform is None


class TestAudioEvalResult:
    def test_add_example(self):
        x = _make_audio_with_text("你好")
        y = _make_audio_with_text("你好")
        result = AudioEvalResult()
        result.add_example(x, y)
        assert len(result) == 1

    def test_add_example_list(self):
        from fasr.data.audio import AudioList

        x_list = AudioList(
            [_make_audio_with_text("你好"), _make_audio_with_text("世界")]
        )
        y_list = AudioList(
            [_make_audio_with_text("你好"), _make_audio_with_text("世界")]
        )
        result = AudioEvalResult()
        result.add_example(x_list, y_list)
        assert len(result) == 2

    def test_add_example_mismatched_raises(self):
        from fasr.data.audio import AudioList

        x_list = AudioList([_make_audio_with_text("你好")])
        y_list = AudioList(
            [_make_audio_with_text("你好"), _make_audio_with_text("世界")]
        )
        result = AudioEvalResult()
        with pytest.raises(ValueError):
            result.add_example(x_list, y_list)

    def test_char_error_rate(self):
        result = AudioEvalResult()
        result.add_example(
            _make_audio_with_text("你好世界"),
            _make_audio_with_text("你好世界"),
        )
        result.add_example(
            _make_audio_with_text("你好世界"),
            _make_audio_with_text("你好世界"),
        )
        cer = result.char_error_rate()
        assert cer == 0.0

    def test_x_y_properties(self):
        result = AudioEvalResult()
        result.add_example(
            _make_audio_with_text("你好"),
            _make_audio_with_text("世界"),
        )
        assert len(result.x) == 1
        assert len(result.y) == 1

    def test_get_badcase(self):
        result = AudioEvalResult()
        result.add_example(
            _make_audio_with_text("你好世界"),
            _make_audio_with_text("你好世界"),
        )
        result.add_example(
            _make_audio_with_text("完全不同"),
            _make_audio_with_text("你好世界"),
        )
        result.char_error_rate()
        badcases = result.get_badcase()
        assert len(badcases) == 1
