import math
from typing import Any, Dict, List, Optional

import numpy as np
import torch
from torch.utils.data import IterableDataset

from lhotse import RecordingSet, Seconds, compute_num_samples, validate
from lhotse.audio.backend import torchaudio_supports_ffmpeg
from lhotse.audio.utils import suppress_audio_loading_errors
from lhotse.augmentation import AugmentFn
from lhotse.cut import CutSet
from lhotse.dataset.collation import collate_audio, collate_features, collate_matrices
from lhotse.features import FeatureExtractor
from lhotse.features.kaldi.layers import _get_strided_batch_streaming


class UnsupervisedDataset(torch.utils.data.Dataset):
    """
    Dataset that contains no supervision - it only provides the features extracted from recordings.

    .. code-block::

        {
            'features': (B x T x F) tensor
            'features_lens': (B, ) tensor
        }
    """

    def __init__(self) -> None:
        super().__init__()

    def __getitem__(self, cuts: CutSet) -> Dict[str, Any]:
        self._validate(cuts)
        features, features_lens = collate_features(cuts)
        return {
            "cuts": cuts,
            "features": features,
            "features_lens": features_lens,
        }

    def _validate(self, cuts: CutSet) -> None:
        validate(cuts)
        assert all(cut.has_features for cut in cuts)


class UnsupervisedWaveformDataset(UnsupervisedDataset):
    """
    A variant of UnsupervisedDataset that provides waveform samples instead of features.
    The output is a tensor of shape (C, T), with C being the number of channels and T the number of audio samples.
    In this implementation, there will always be a single channel.

    Returns:

    .. code-block::

        {
            'audio': (B x NumSamples) float tensor
            'audio_lens': (B, ) int tensor
        }
    """

    def __init__(self, collate: bool = True) -> None:
        super().__init__()
        self.collate = collate

    def __getitem__(self, cuts: CutSet) -> Dict[str, Any]:
        if self.collate:
            audio, audio_lens = collate_audio(cuts)
            return {
                "cuts": cuts,
                "audio": audio,
                "audio_lens": audio_lens,
            }
        else:
            remain_cuts = []
            remain_audios = []
            for c in cuts:
                with suppress_audio_loading_errors():
                    remain_audios.append(c.load_audio())
                    remain_cuts.append(c)
            return {"cuts": CutSet.from_cuts(remain_cuts), "audio": remain_audios}

    def _validate(self, cuts: CutSet) -> None:
        validate(cuts)
        assert all(cut.has_recording for cut in cuts)


class DynamicUnsupervisedDataset(UnsupervisedDataset):
    """
    An example dataset that shows how to use on-the-fly feature extraction in Lhotse.
    It accepts two additional inputs - a FeatureExtractor and an optional WavAugmenter for time-domain data augmentation..
    The output is approximately the same as that of the ``UnsupervisedDataset`` -
    there might be slight differences for ``MixedCut``s, because this dataset mixes them in the time domain,
    and ``UnsupervisedDataset`` does that in the feature domain.
    Cuts that are not mixed will yield identical results in both dataset classes.
    """

    def __init__(
        self,
        feature_extractor: FeatureExtractor,
        augment_fn: Optional[AugmentFn] = None,
    ):
        super().__init__()
        self.feature_extractor = feature_extractor
        self.augment_fn = augment_fn

    def __getitem__(self, cuts: CutSet) -> torch.Tensor:
        self._validate(cuts)

        def generate_cut(cuts: CutSet):
            for cut in cuts:
                with suppress_audio_loading_errors():
                    yield cut.compute_features(
                        extractor=self.feature_extractor,
                        augment_fn=self.augment_fn,
                    )

        features = collate_matrices(generate_cut(cuts))
        return features

    def _validate(self, cuts: CutSet) -> None:
        validate(cuts)
        assert all(cut.has_recording for cut in cuts)


class RecordingChunkIterableDataset(IterableDataset):
    """
    This dataset iterates over chunks of a recording, for each recording provided.
    It supports setting a chunk_shift < chunk_size to run model predictions on
    overlapping audio chunks.

    The format of yielded items is the following::

        {
            "recording_id": str
            "begin_time": tensor with dtype=float32 shape=(1,)
            "end_time": tensor with dtype=float32 shape=(1,)
            "audio": tensor with dtype=float32 shape=(chunk_size_in_samples,)
        }

    Unlike most other datasets in Lhotse, this dataset does not yield batched items,
    and should be used like the following::

        >>> recordings = RecordingSet.from_file("my-recordings.jsonl.gz")
        ... dataset = RecordingChunkIterableDataset(recordings, chunk_size=30.0, chunk_shift=25.0)
        ... dloader = torch.utils.data.DataLoader(
        ...     dataset,
        ...     batch_size=32,
        ...     collate_fn=audio_chunk_collate,
        ...     worker_init_fn=audio_chunk_worker_init_fn,
        ... )

    """

    def __init__(
        self, recordings: RecordingSet, chunk_size: Seconds, chunk_shift: Seconds
    ) -> None:
        self.recordings = list(recordings)
        self.chunk_size = chunk_size
        self.chunk_shift = chunk_shift
        self.start = 0
        self.end = len(self.recordings)

        self.validate()

    def validate(self) -> None:
        for r in self.recordings:
            assert (
                len(r.sources) == 1
            ), f"We currently don't support multi-source audio in this dataset (got {len(r.sources)} sources in recording {r.id})."
            assert (
                r.sources[0].type == "file"
            ), f"We currently only support 'file' AudioSource type in this dataset (got: {r.sources[0].type} in recording {r.id})."
            assert (
                r.num_channels == 1
            ), f"We currently only support single-channel audio in this dataset (got {r.num_channels} channels in recording {r.id})."

    def __iter__(self):
        import soundfile as sf

        for r in self.recordings[self.start : self.end]:
            chunk_size = compute_num_samples(self.chunk_size, r.sampling_rate)
            chunk_overlap = compute_num_samples(
                self.chunk_size - self.chunk_shift, r.sampling_rate
            )

            begin_time = 0
            end_time = self.chunk_size
            with sf.SoundFile(r.sources[0].source, "rb") as stream:
                for chunk in stream.blocks(
                    chunk_size, overlap=chunk_overlap, dtype=np.float32
                ):
                    chunk = torch.as_tensor(chunk)
                    yield {
                        "recording_id": r.id,
                        "begin_time": torch.as_tensor(begin_time, dtype=torch.float32),
                        "end_time": torch.as_tensor(end_time, dtype=torch.float32),
                        "audio": chunk,
                    }
                    begin_time += self.chunk_shift
                    end_time = begin_time + self.chunk_size


def audio_chunk_collate(batch: List[Dict]):
    from torch.utils.data import default_collate

    audios = [d.pop("audio") for d in batch]
    out = default_collate(batch)

    maxlen = max(a.shape[0] for a in audios)
    audio = torch.zeros((len(audios), maxlen))
    for i, a in enumerate(audios):
        audio[i, : a.shape[0]] = torch.as_tensor(a)
    out["audio"] = audio

    return out


def audio_chunk_worker_init_fn(worker_id):
    worker_info = torch.utils.data.get_worker_info()
    dataset = worker_info.dataset  # the dataset copy in this worker process
    overall_start = dataset.start
    overall_end = dataset.end
    # configure the dataset to only process the split workload
    per_worker = int(
        math.ceil((overall_end - overall_start) / float(worker_info.num_workers))
    )
    worker_id = worker_info.id
    dataset.start = overall_start + worker_id * per_worker
    dataset.end = min(dataset.start + per_worker, overall_end)
