import math
import random
import uuid
from contextlib import AbstractContextManager, contextmanager
from dataclasses import asdict, dataclass
from decimal import Decimal, ROUND_HALF_DOWN, ROUND_HALF_UP
from math import ceil, isclose
from pathlib import Path
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, TypeVar, Union

import numpy as np
import torch
from tqdm.auto import tqdm

Pathlike = Union[Path, str]
T = TypeVar('T')

Seconds = float
Decibels = float

INT16MAX = 32768
EPSILON = 1e-10
LOG_EPSILON = math.log(EPSILON)

# This is a utility that generates uuid4's and is set when the user calls
# the ``fix_random_seed`` function.
# Python's uuid module is not affected by the ``random.seed(value)`` call,
# so we work around it to provide deterministic ID generation when requested.
_lhotse_uuid: Optional[Callable] = None


def fix_random_seed(random_seed: int):
    """
    Set the same random seed for the libraries and modules that Lhotse interacts with.
    Includes the ``random`` module, numpy, torch, and ``uuid4()`` function defined in this file.
    """
    global _lhotse_uuid
    random.seed(random_seed)
    np.random.seed(random_seed)
    torch.random.manual_seed(random_seed)
    # Ensure deterministic ID creation
    rd = random.Random()
    rd.seed(random_seed)
    _lhotse_uuid = lambda: uuid.UUID(int=rd.getrandbits(128))


def uuid4():
    """
    Generates uuid4's exactly like Python's uuid.uuid4() function.
    When ``fix_random_seed()`` is called, it will instead generate deterministic IDs.
    """
    if _lhotse_uuid is not None:
        return _lhotse_uuid()
    return uuid.uuid4()


def asdict_nonull(dclass) -> Dict[str, Any]:
    """
    Recursively convert a dataclass into a dict, removing all the fields with `None` value.
    Intended to use in place of dataclasses.asdict(), when the null values are not desired in the serialized document.
    """
    def non_null_dict_factory(collection):
        d = dict(collection)
        remove_keys = []
        for key, val in d.items():
            if val is None:
                remove_keys.append(key)
        for k in remove_keys:
            del d[k]
        return d

    return asdict(dclass, dict_factory=non_null_dict_factory)




class SetContainingAnything:
    def __contains__(self, item):
        return True

    def intersection(self, iterable):
        return True


@dataclass
class TimeSpan:
    """Helper class for specifying a time span."""
    start: Seconds
    end: Seconds


# TODO: Ugh, Protocols are only in Python 3.8+...
def overlaps(lhs: Any, rhs: Any) -> bool:
    """Indicates whether two time-spans/segments are overlapping or not."""
    return lhs.start < rhs.end and rhs.start < lhs.end \
           and not isclose(lhs.start, rhs.end) and not isclose(rhs.start, lhs.end)


def overspans(spanning: Any, spanned: Any) -> bool:
    """Indicates whether the left-hand-side time-span/segment covers the whole right-hand-side time-span/segment."""
    return spanning.start <= spanned.start <= spanned.end <= spanning.end


def time_diff_to_num_frames(time_diff: Seconds, frame_length: Seconds, frame_shift: Seconds) -> int:
    """Convert duration to an equivalent number of frames, so as to not exceed the duration."""
    if isclose(time_diff, 0.0):
        return 0
    return int(ceil((time_diff - frame_length) / frame_shift))


def check_and_rglob(path: Pathlike, pattern: str, strict: Optional[bool] = True) -> List[Path]:
    """
    Asserts that ``path`` exists, is a directory and contains at least one file satisfying the ``pattern``.
    If `strict` is False, then zero matches are allowed.

    :returns: a list of paths to files matching the ``pattern``.
    """
    path = Path(path)
    assert path.is_dir(), f'No such directory: {path}'
    matches = sorted(path.rglob(pattern))
    assert len(matches) > 0 or not strict, f'No files matching pattern "{pattern}" in directory: {path}'
    return matches


@contextmanager
def recursion_limit(stack_size: int):
    """
    Code executed in this context will be able to recurse up to the specified recursion limit
    (or will hit a StackOverflow error if that number is too high).

    Usage:
        >>> with recursion_limit(1000):
        >>>     pass
    """
    import sys
    old_size = sys.getrecursionlimit()
    sys.setrecursionlimit(stack_size)
    try:
        yield
    finally:
        sys.setrecursionlimit(old_size)


def fastcopy(dataclass_obj: T, **kwargs) -> T:
    """
    Returns a new object with the same member values.
    Selected members can be overwritten with kwargs.
    It's supposed to work only with dataclasses.
    It's 10X faster than the other methods I've tried...

    Example:
        >>> ts1 = TimeSpan(start=5, end=10)
        >>> ts2 = fastcopy(ts1, end=12)
    """
    return type(dataclass_obj)(**{**dataclass_obj.__dict__, **kwargs})


def split_sequence(
        seq: Sequence[Any],
        num_splits: int,
        shuffle: bool = False,
        drop_last: bool = False
) -> List[List[Any]]:
    """
    Split a sequence into ``num_splits`` equal parts. The element order can be randomized.
    Raises a ``ValueError`` if ``num_splits`` is larger than ``len(seq)``.

    :param seq: an input iterable (can be a Lhotse manifest).
    :param num_splits: how many output splits should be created.
    :param shuffle: optionally shuffle the sequence before splitting.
    :param drop_last: determines how to handle splitting when ``len(seq)`` is not divisible
        by ``num_splits``. When ``False`` (default), the splits might have unequal lengths.
        When ``True``, it may discard the last element in some splits to ensure they are
        equally long.
    :return: a list of length ``num_splits`` containing smaller lists (the splits).
    """
    seq = list(seq)
    num_items = len(seq)
    if num_splits > num_items:
        raise ValueError(f"Cannot split iterable into more chunks ({num_splits}) than its number of items {num_items}")
    if shuffle:
        random.shuffle(seq)
    chunk_size = num_items // num_splits

    num_shifts = num_items % num_splits
    if drop_last:
        # Equally-sized splits; discards the remainder by default, no shifts are needed
        end_shifts = [0] * num_splits
        begin_shifts = [0] * num_splits
    else:
        # Non-equally sized splits; need to shift the indices like:
        # [0, 10] -> [0, 11]    (begin_shift=0, end_shift=1)
        # [10, 20] -> [11, 22]  (begin_shift=1, end_shift=2)
        # [20, 30] -> [22, 32]  (begin_shift=2, end_shift=2)
        # for num_items=32 and num_splits=3
        end_shifts = list(range(1, num_shifts + 1)) + [num_shifts] * (num_splits - num_shifts)
        begin_shifts = [0] + end_shifts[:-1]

    split_indices = [
        [i * chunk_size + begin_shift, (i + 1) * chunk_size + end_shift]
        for i, begin_shift, end_shift in zip(range(num_splits), begin_shifts, end_shifts)
    ]
    splits = [seq[begin:end] for begin, end in split_indices]
    return splits


def compute_num_frames(
        duration: Seconds,
        frame_shift: Seconds,
        sampling_rate: int,
) -> int:
    """
    Compute the number of frames from duration and frame_shift in a safe way.
    """
    num_samples = round(duration * sampling_rate)
    window_hop = round(frame_shift * sampling_rate)
    num_frames = int((num_samples + window_hop // 2) // window_hop)
    return num_frames


def during_docs_build() -> bool:
    import os
    return bool(os.environ.get('READTHEDOCS'))


def tqdm_urlretrieve_hook(t):
    """Wraps tqdm instance.
    Don't forget to close() or __exit__()
    the tqdm instance once you're done with it (easiest using `with` syntax).
    Example
    -------
    >>> from urllib.request import urlretrieve
    >>> with tqdm(...) as t:
    ...     reporthook = tqdm_urlretrieve_hook(t)
    ...     urlretrieve(..., reporthook=reporthook)

    Source: https://github.com/tqdm/tqdm/blob/master/examples/tqdm_wget.py
    """
    last_b = [0]

    def update_to(b=1, bsize=1, tsize=None):
        """
        b  : int, optional
            Number of blocks transferred so far [default: 1].
        bsize  : int, optional
            Size of each block (in tqdm units) [default: 1].
        tsize  : int, optional
            Total size (in tqdm units). If [default: None] or -1,
            remains unchanged.
        """
        if tsize not in (None, -1):
            t.total = tsize
        displayed = t.update((b - last_b[0]) * bsize)
        last_b[0] = b
        return displayed

    return update_to


def urlretrieve_progress(url, filename=None, data=None, desc=None):
    """
    Works exactly like urllib.request.urlretrieve, but attaches a tqdm hook to display
    a progress bar of the download.
    Use "desc" argument to display a user-readable string that informs what is being downloaded.
    """
    from urllib.request import urlretrieve
    with tqdm(unit='B', unit_scale=True, unit_divisor=1024, miniters=1, desc=desc) as t:
        reporthook = tqdm_urlretrieve_hook(t)
        return urlretrieve(url=url, filename=filename, reporthook=reporthook, data=data)


class nullcontext(AbstractContextManager):
    """Context manager that does no additional processing.

    Used as a stand-in for a normal context manager, when a particular
    block of code is only sometimes used with a normal context manager:

    cm = optional_cm if condition else nullcontext()
    with cm:
        # Perform operation, using optional_cm if condition is True

    Note(pzelasko): This is copied from Python 3.7 stdlib so that we can use it in 3.6.
    """

    def __init__(self, enter_result=None):
        self.enter_result = enter_result

    def __enter__(self):
        return self.enter_result

    def __exit__(self, *excinfo):
        pass


def perturb_num_samples(num_samples: int, factor: float) -> int:
    """Mimicks the behavior of the speed perturbation on the number of samples."""
    rounding = ROUND_HALF_UP if factor >= 1.0 else ROUND_HALF_DOWN
    return int(Decimal(round(num_samples / factor, ndigits=8)).quantize(0, rounding=rounding))


def compute_num_samples(duration: Seconds, sampling_rate: int, rounding=ROUND_HALF_UP) -> int:
    """
    Convert a time quantity to the number of samples given a specific sampling rate.
    Performs consistent rounding up or down for ``duration`` that is not a multiply of
    the sampling interval (unlike Python's built-in ``round()`` that implements banker's rounding).
    """
    return int(
        Decimal(
            round(duration * sampling_rate, ndigits=8)
        ).quantize(0, rounding=rounding)
    )


def index_by_id_and_check(manifests: Iterable[T]) -> Dict[str, T]:
    id2man = {}
    for m in manifests:
        assert m.id not in id2man, f"Duplicated manifest ID: {m.id}"
        id2man[m.id] = m
    return id2man


def exactly_one_not_null(*args) -> bool:
    not_null = [arg is not None for arg in args]
    return sum(not_null) == 1


def supervision_to_frames(
        supervision,
        frame_shift: Seconds,
        sampling_rate: int,
        max_frames: Optional[int] = None
) -> Tuple[int, int]:
    """
    Utility to convert a supervision's time span into a tuple of ``(start_frame, num_frames)``.
    When ``max_frames`` is specified, it will truncate the ``num_frames`` (if necessary).
    """
    start_frame = compute_num_frames(
        supervision.start,
        frame_shift=frame_shift,
        sampling_rate=sampling_rate
    )
    num_frames = compute_num_frames(
        supervision.duration,
        frame_shift=frame_shift,
        sampling_rate=sampling_rate
    )
    if max_frames:
        diff = start_frame + num_frames - max_frames
        if diff > 0:
            num_frames -= diff
    return start_frame, num_frames


def supervision_to_samples(
        supervision,
        sampling_rate: int,
        max_samples: Optional[int] = None
) -> Tuple[int, int]:
    """
    Utility to convert a supervision's time span into a tuple of ``(start_sample num_samples)``.
    When ``max_samples`` is specified, it will truncate the ``num_samples`` (if necessary).
    """
    start_sample = compute_num_samples(
        supervision.start,
        sampling_rate=sampling_rate
    )
    num_samples = compute_num_samples(
        supervision.duration,
        sampling_rate=sampling_rate
    )
    if max_samples:
        diff = start_sample + num_samples - max_samples
        if diff > 0:
            num_samples -= diff
    return start_sample, num_samples


def is_none_or_gt(value, threshold) -> bool:
    return value is None or value > threshold


def is_module_available(*modules: str) -> bool:
    r"""Returns if a top-level module with :attr:`name` exists *without**
    importing it. This is generally safer than try-catch block around a
    `import X`. It avoids third party libraries breaking assumptions of some of
    our tests, e.g., setting multiprocessing start method when imported
    (see librosa/#747, torchvision/#544).

    Note: "borrowed" from torchaudio:
    https://github.com/pytorch/audio/blob/6bad3a66a7a1c7cc05755e9ee5931b7391d2b94c/torchaudio/_internal/module_utils.py#L9
    """
    import importlib
    return all(importlib.util.find_spec(m) is not None for m in modules)


def measure_overlap(lhs: Any, rhs: Any) -> float:
    """
    Given two objects with "start" and "end" attributes, return the % of their overlapped time
    with regard to the shorter of the two spans.
    ."""
    lhs, rhs = sorted([lhs, rhs], key=lambda item: item.start)
    overlapped_area = lhs.end - rhs.start
    if overlapped_area <= 0:
        return 0.
    dur = min(lhs.end - lhs.start, rhs.end - rhs.start)
    return overlapped_area / dur


def ifnone(item: Optional[Any], alt_item: Any) -> Any:
    """Return ``alt_item`` if ``item is None``, otherwise ``item``."""
    return alt_item if item is None else item


def lens_to_mask(lens: torch.IntTensor) -> torch.Tensor:
    """
    Create a 2-D mask tensor of shape (batch_size, max_length) and dtype float32
    from a 1-D tensor of integers describing the length of batch samples in another tensor.
    """
    mask = lens.new_zeros(lens.shape[0], max(lens), dtype=torch.float32)
    for i, num in enumerate(lens):
        mask[i, :num] = 1.0
    return mask


class NonPositiveEnergyError(ValueError):
    pass
