import tensorflow as tf
from tensorflow import data as tfdata

from audioclass.batch.base import BaseIterator, BatchGenerator
from audioclass.preprocess import load_recording

__all__ = [
    "TFDatasetIterator",
]


class TFDatasetIterator(BaseIterator):
    def __iter__(self) -> BatchGenerator:
        @tf.py_function(Tout=tf.float32)  # type: ignore
        def apply_fn(index):  # pragma: no cover
            index = index.numpy()
            recording = self.recordings[index]
            return load_recording(
                recording,
                samplerate=self.samplerate,
                buffer_size=self.input_samples,
                audio_dir=self.audio_dir,
            )

        dataset = (
            tfdata.Dataset.range(len(self.recordings))
            .map(
                apply_fn,
                name="load_recording",
                num_parallel_calls=tf.data.AUTOTUNE,
            )
            .enumerate()
            .flat_map(
                lambda ind, arr: tfdata.Dataset.zip(
                    tfdata.Dataset.from_tensor_slices(arr, name="audio"),
                    tfdata.Dataset.from_tensors([ind]).repeat().enumerate(),
                    name="zip",
                )
            )
            .batch(self.batch_size, drop_remainder=False)
            .prefetch(tf.data.AUTOTUNE)
        )

        for batch_data, (frame, index) in dataset.as_numpy_iterator():  # type: ignore
            yield batch_data, [self.recordings[int(i)] for i in index], frame
