"""
:class:`ConcatFilesDataset `

https://github.com/rwth-i6/returnn/issues/1519
"""

from __future__ import annotations

from typing import Optional, Any, Dict, Tuple, List, Sequence, Collection, Callable, Union
import os
import sys
import numpy
from returnn.log import log
from returnn.util.basic import try_run
from returnn.util.multi_proc_non_daemonic_spawn import NonDaemonicSpawnContext
from returnn.config import SubProcCopyGlobalConfigPreInitFunc
from .basic import init_dataset, DatasetSeq
from .cached2 import CachedDataset2

# noinspection PyProtectedMember
from multiprocessing.connection import Connection as mpConnection

_mp = NonDaemonicSpawnContext(process_pre_init_func=SubProcCopyGlobalConfigPreInitFunc())


__all__ = ["ConcatFilesDataset"]

Filename = str
FileTree = Union[Filename, Tuple["FileTree", ...], Dict[Any, "FileTree"], List["FileTree"]]


class ConcatFilesDataset(CachedDataset2):
    """
    This is similar to :class:`ConcatDataset`, but instead of concatenating datasets,
    we distribute files over subepochs,
    and then create a sub dataset for every sub epoch
    for a given subset of the files
    via the given ``get_sub_epoch_dataset``.

    This scheme allows to shuffle over the files,
    which makes shuffling much more efficient over a large dataset
    at the cost of no longer shuffling over the full dataset in every subepoch.
    Instead, the quality of the shuffle depends on the number of files the dataset is
    split into -- the more files per subepoch, the better.

    Additionally, this scheme allows to prefetch and cache the upcoming needed files,
    e.g. copying them from a NFS to local disk,
    when the local disk can not store the whole dataset,
    and/or when the local disk would allow for faster access.
    For that, the user can use :class:`returnn.util.file_cache.CachedFile`
    inside the returned dict from ``get_sub_epoch_dataset``.

    It was also designed with multi-GPU training in mind,
    where each GPU should get a set of files to work on,
    and the sub epochs across the GPUs would not be exactly of the same length.
    Whenever one GPU finishes its sub epoch, all remaining data on the other GPUs is dropped.
    First, our scheme tries to make the length of the sub epochs as equally distributed as possible
    after random shuffling by using the size of the files as proxy for the length of the contents.
    Further, due to the random shuffling, there should not be any bias in the data distribution.
    Specifically, we don't want that some data might be visited more often than others
    (at least its expected value should be the same).

    Example usage::

        def get_sub_epoch_dataset(files_subepoch: List[str]) -> Dict[str, Any]:
          from returnn.util.file_cache import CachedFile
          return {
            "class": "HDFDataset",
            "files": [CachedFile(fn) for fn in files_subepoch],
          }

        train = {
          "class": "ConcatFilesDataset",
          "files": [
            "/nfs/big_data_1.hdf",
            ...
          ],  # M files
          "get_sub_epoch_dataset": get_sub_epoch_dataset,
          "partition_epoch": P,  # P << M
        }

    Instead of a plain list of strings, you can also provide a list of any nested structures to keep
    multimodal data together::

        def get_sub_epoch_dataset(files_subepoch: List[Tuple[str, str]]) -> Dict[str, Any]:
          from returnn.util.file_cache import CachedFile

          alignments, features = tuple(zip(*files_subepoch)) # transpose

          return {
            "class": "MetaDataset",
            "data_map": {"classes": ("alignments", "data"), "data": ("features", "data")},
            "datasets": {
              "alignments": {
                "class": "HDFDataset",
                "files": alignments,
                "seq_ordering": "random",
              },
              "features": {
                "class": "HDFDataset",
                "files": features,
              },
            },
            "seq_order_control_dataset": "alignments",
          }

        train = {
          "class": "ConcatFilesDataset",
          "files": [
            ("/nfs/alignment_1.hdf", "/nfs/features_1.hdf"),
            ...
          ],  # M entries
          "get_sub_epoch_dataset": get_sub_epoch_dataset,
          "partition_epoch": P,  # P << M
        }

    In this case the file sizes for sub epoch distribution are summed up per list entry
    by iterating over the structure leaves.

    For some discussion, see https://github.com/rwth-i6/returnn/issues/1519 and
    https://github.com/rwth-i6/returnn/issues/1524.
    """

    def __init__(
        self,
        *,
        files: List[FileTree],
        get_sub_epoch_dataset: Callable[[List[FileTree]], Dict[str, Any]],
        preload_next_n_sub_epochs: int = 1,
        buffer_size: int = 1,
        file_cache_opts: Optional[Dict[str, Any]] = None,
        _meta_info_cache: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        """
        :param files: the files to shuffle over, can also be a list of arbitrarily nested python objects
            to keep associated heterogeneous data together
        :param get_sub_epoch_dataset: callable which returns a dataset dict for a given subset of files
        :param preload_next_n_sub_epochs: how many sub epoch datasets to preload
        :param buffer_size: buffer size for each worker, amount of seqs to prefetch
        :param _meta_info_cache: for internal use
        """
        super().__init__(**kwargs)
        self.files = files
        self.get_sub_epoch_dataset = get_sub_epoch_dataset
        assert preload_next_n_sub_epochs >= 0
        self.preload_next_n_sub_epochs = preload_next_n_sub_epochs
        self.buffer_size = buffer_size
        self.file_cache_opts = file_cache_opts or {}
        self._file_sizes: Optional[Dict[str, int]] = None  # key -> size. for equal distribution across sub epochs
        self._data_keys: Optional[List[str]] = None
        self._num_seqs: Optional[int] = None

        self._file_cache: Optional[_FileCacheProc] = None
        self._workers: Dict[int, _WorkerProcParent] = {}  # epoch -> worker
        self._files_order_cache: Dict[int, List[List[FileTree]]] = {}  # full epoch (0-indexed) -> files order

        if _meta_info_cache:
            # This allows to skip the lazy init in self.initialize().
            # This can be used when pickling/deepcopying an instance of this dataset,
            # see Dataset.__reduce__.
            self.num_inputs = _meta_info_cache["num_inputs"]
            self.num_outputs = _meta_info_cache["num_outputs"]
            self.labels = _meta_info_cache["labels"]
            self._data_keys = _meta_info_cache["data_keys"]
            self._file_sizes = _meta_info_cache["file_sizes"]

        if len(files) < self.partition_epoch:
            raise ValueError(f"{self}: len(files) {len(files)} < partition_epoch {self.partition_epoch}")

    def initialize(self):
        """init"""
        self._lazy_init_num_outputs()
        super().initialize()

    @property
    def _meta_info_cache(self):
        if not self.num_outputs:
            return None
        return {
            "num_inputs": self.num_inputs,
            "num_outputs": self.num_outputs,
            "labels": self.labels,
            "data_keys": self._data_keys,
            "file_sizes": self._file_sizes,
        }

    def _lazy_init_num_outputs(self):
        if self.num_outputs:
            return
        # First, we need to know the num_inputs, num_outputs, total_num_seqs, labels.
        # Init the dataset with the first file.
        dataset_dict, exit_hook = self._get_sub_dataset_dict(files=[self.files[0]])
        try:
            dataset = init_dataset(dataset_dict, extra_kwargs={"seq_ordering": "default"})
            self.num_inputs = dataset.num_inputs
            self.num_outputs = dataset.num_outputs
            self.labels = dataset.labels
            self._data_keys = dataset.get_data_keys()
        finally:
            exit_hook()

    def _lazy_init_file_sizes(self):
        import tree

        if self._file_sizes:
            return
        self._file_sizes = {
            _get_key_for_file_tree(t): sum((os.path.getsize(fn) for fn in tree.flatten(t)), 0) for t in self.files
        }

    def _lazy_init_file_cache_proc(self):
        if self._file_cache:
            return
        self._file_cache = _FileCacheProc(
            name=f"{self.__class__.__name__} {self.name}", file_cache_opts=self.file_cache_opts
        )

    def __del__(self):
        if self._file_cache:
            try_run(self._file_cache.exit, kwargs={"join": False})
            self._file_cache = None
        for k, worker in self._workers.items():
            try_run(worker.exit, kwargs={"join": False})
        self._workers.clear()

    def init_seq_order(self, epoch: Optional[int] = None, seq_list=None, seq_order=None) -> bool:
        """
        :param epoch:
        :param seq_list:
        :param seq_order:
        :return: whether the order changed (True is always safe to return)
        """
        super().init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
        if seq_list is not None:
            raise Exception(f"{self}: seq_list not supported")
        if seq_order is not None:
            raise Exception(f"{self}: seq_order not supported")
        if epoch is None:
            self._num_seqs = 0
            return True

        self._lazy_init_file_sizes()
        self._lazy_init_file_cache_proc()

        # Cleanup and fill _files_order_cache.
        full_epoch_0idx = (epoch - 1) // self.partition_epoch
        for k in list(self._files_order_cache.keys()):
            if k < full_epoch_0idx:
                del self._files_order_cache[k]
        for ep_ in range(epoch, epoch + self.preload_next_n_sub_epochs + 1):
            full_epoch_0idx_ = (ep_ - 1) // self.partition_epoch
            if full_epoch_0idx_ in self._files_order_cache:
                continue
            if self.seq_ordering == "default":
                files_order_flat = self.files
            elif self.seq_ordering == "random":
                rnd_seed = self._get_random_seed_for_epoch(full_epoch_0idx_ * self.partition_epoch + 1)
                random_generator = numpy.random.RandomState(rnd_seed)
                files_order_flat = list(self.files)
                random_generator.shuffle(files_order_flat)
            else:
                raise ValueError(f"{self}: seq_ordering {self.seq_ordering!r} not supported")
            self._files_order_cache[full_epoch_0idx_] = self._get_files_per_sub_epochs(
                partition_epoch=self.partition_epoch, file_sizes=self._file_sizes, files_order=files_order_flat
            )

        # Cleanup and fill _workers.
        for k, worker in list(self._workers.items()):
            if k < epoch:
                worker.exit()
                del self._workers[k]
        for ep_ in range(epoch, epoch + self.preload_next_n_sub_epochs + 1):
            if ep_ in self._workers:
                continue
            full_epoch_0idx_ = (ep_ - 1) // self.partition_epoch
            files_order: List[List[FileTree]] = self._files_order_cache[full_epoch_0idx_]
            files_for_subep = files_order[(ep_ - 1) % self.partition_epoch]
            print(f"{self}: using files for epoch {ep_}: {files_for_subep}", file=log.v4)
            dataset_dict, exit_hook = self._get_sub_dataset_dict(files=files_for_subep)
            worker = _WorkerProcParent(
                name=f"{self.__class__.__name__} {self.name} ep {epoch}",
                epoch=ep_,
                full_epoch_0idx=full_epoch_0idx_,
                dataset_dict=dataset_dict,
                buffer_size=self.buffer_size,
                exit_hook=exit_hook,
            )
            self._workers[ep_] = worker

        self._num_seqs = self._workers[epoch].get_num_seqs()
        return True

    def _get_sub_dataset_dict(self, files: List[FileTree]) -> Tuple[Dict[str, Any], _FileCacheExitHook]:
        dataset_dict = self.get_sub_epoch_dataset(files)
        if "random_seed_offset" not in dataset_dict:
            dataset_dict["random_seed_offset"] = self.random_seed_offset
        if dataset_dict.get("partition_epoch", 1) != 1:
            raise ValueError(f"{self}: sub dataset should not have partition_epoch, got: {dataset_dict}")
        if "seq_ordering" not in dataset_dict and "seq_order_control_dataset" not in dataset_dict:
            raise ValueError(
                f"{self}: sub dataset should have explicit seq_ordering "
                "(or seq_order_control_dataset for MetaDataset), got: {dataset_dict}"
            )
        self._lazy_init_file_cache_proc()
        dataset_dict, exit_hook = self._file_cache.handle_cached_files_in_config(dataset_dict)
        return dataset_dict, exit_hook

    @staticmethod
    def _get_files_per_sub_epochs(
        *, partition_epoch: int, file_sizes: Dict[str, int], files_order: Sequence[FileTree]
    ) -> List[List[FileTree]]:
        total_size = sum(file_sizes.values())
        avg_size_per_sub_epoch = total_size / partition_epoch
        # Now evenly distribute the files over the sub epochs.
        # Note that many one-pass variants of algorithms to evenly distribute
        # can end up with some empty sub epochs,
        # so we need to make sure that this is not the case.
        # E.g. consider the seqs of size [1,1,78,120] and partition_epoch=4.
        # That has avg size per sub epoch 50.
        # Some simple algorithms could end up with the sub epochs
        # [[1,1], [78], [120], []] or [[1,1,78], [120], [], []].
        # Or consider [5,5]+[10]*7, partition_epoch=5, which has avg size 16.
        # A simple algorithm could end up with [[5,5,10], [10,10], [10,10], [10,10], []].
        # See test_ConcatFilesDataset_get_files_per_sub_epochs for some test cases.
        assert len(files_order) >= partition_epoch
        files_per_sub_epochs = [[] for _ in range(partition_epoch)]
        assert len(files_per_sub_epochs) == partition_epoch
        sub_epoch_idx = 0
        size_taken = 0
        for i, f_tree in enumerate(files_order):
            size = file_sizes[_get_key_for_file_tree(f_tree)]
            num_remaining = len(files_order) - i
            if num_remaining <= partition_epoch - sub_epoch_idx - 1:
                # All remaining sub epochs must be filled.
                assert files_per_sub_epochs[sub_epoch_idx]
                sub_epoch_idx += 1
                files_per_sub_epochs[sub_epoch_idx].append(f_tree)
                size_taken = size
                continue
            if sub_epoch_idx == partition_epoch - 1:
                # We are done. Just add the rest to the last sub epoch.
                files_per_sub_epochs[sub_epoch_idx].append(f_tree)
                size_taken += size
                continue
            assert size_taken <= avg_size_per_sub_epoch
            if size_taken + size <= avg_size_per_sub_epoch:
                files_per_sub_epochs[sub_epoch_idx].append(f_tree)
                size_taken += size
                continue
            # We should increase the sub epoch index.
            # We need to decide where to add this file, to the current or the next sub epoch.
            if not files_per_sub_epochs[sub_epoch_idx] or (
                # Better to add this file to the current sub epoch?
                abs((size_taken + size) - avg_size_per_sub_epoch)
                <= abs(size_taken - avg_size_per_sub_epoch)
            ):
                files_per_sub_epochs[sub_epoch_idx].append(f_tree)
                size_taken = 0
            else:
                files_per_sub_epochs[sub_epoch_idx + 1].append(f_tree)
                size_taken = size
            sub_epoch_idx += 1
        assert all(files_per_sub_epochs)
        return files_per_sub_epochs

    def _collect_single_seq(self, seq_idx: int) -> Optional[DatasetSeq]:
        if seq_idx >= self._num_seqs:
            return None
        seq = self._workers[self.epoch].get_data_seq(seq_idx)
        if not seq:
            return None
        seq.seq_idx = seq_idx
        return seq

    def have_seqs(self) -> bool:
        """have seqs"""
        return bool(self.files)

    @property
    def num_seqs(self) -> int:
        """num seqs"""
        return self._num_seqs

    def finish_epoch(self, *, free_resources: bool = False):
        """finish epoch"""
        super().finish_epoch(free_resources=free_resources)
        if free_resources:
            for worker in self._workers.values():
                worker.exit()
            self._workers.clear()
            self._files_order_cache.clear()
            if self._file_cache:
                self._file_cache.exit()
                self._file_cache = None
        else:
            if self.epoch in self._workers:
                worker = self._workers.pop(self.epoch)
                worker.exit()

    def get_data_keys(self) -> List[str]:
        """data keys"""
        if self._data_keys is None:
            self._lazy_init_num_outputs()
        return self._data_keys


def _get_key_for_file_tree(t: FileTree) -> str:
    """generates a deterministic key given a file tree"""
    import tree

    return ":".join(tree.flatten(t))


class _WorkerProcParent:
    def __init__(
        self,
        *,
        name: str,
        epoch: int,
        full_epoch_0idx: int,
        dataset_dict: Dict[str, Any],
        buffer_size: int,
        exit_hook: Optional[Callable[[], None]] = None,
    ):
        self.epoch = epoch
        self.full_epoch_0idx = full_epoch_0idx
        self.dataset_dict = dataset_dict
        self.exit_hook = exit_hook

        parent_conn, child_conn = _mp.Pipe()
        self.parent_conn: mpConnection = parent_conn

        self.worker_proc = _mp.Process(
            name=f"{name} worker ep {epoch}",
            target=_worker_proc_loop,
            args=(epoch, buffer_size, dataset_dict, child_conn),
            daemon=True,
        )
        self.worker_proc.start()
        # Make sure the child connection is closed here.
        # It stays open in the child, until the child dies.
        # When that happens, now any consecutive read on the pipe
        # should yield an exception -- which is what we want,
        # otherwise it would just hang.
        child_conn.close()

        # This worker operates on one sub epoch.
        # The data used here is only for this sub epoch.
        # Without shuffling, the next time we would use this data is in the next full epoch.
        # Thus, use the full epoch as the epoch here.
        # Make sure this init is async - we should not wait for this in the main process when it is not used yet.
        self.parent_conn.send(("init_seq_order", {"epoch": self.full_epoch_0idx + 1}))
        self._num_seqs: Optional[int] = None

    def _lazy_wait_for_init_seq_order(self):
        if self._num_seqs is not None:
            return
        msg, num_seqs = self.parent_conn.recv()
        assert msg == "num_seqs" and isinstance(num_seqs, int)
        self._num_seqs = num_seqs

    def get_num_seqs(self) -> int:
        """num seqs for this sub epoch"""
        self._lazy_wait_for_init_seq_order()
        return self._num_seqs

    def get_data_seq(self, seq_idx: int) -> Optional[DatasetSeq]:
        """get data seq"""
        self._lazy_wait_for_init_seq_order()
        self.parent_conn.send(("get_data_seq", {"seq_idx": seq_idx}))
        msg, data = self.parent_conn.recv()
        assert msg == "data_seq"
        return data

    def exit(self, *, join: bool = True):
        """exit"""
        self._lazy_wait_for_init_seq_order()
        self.parent_conn.send(("exit", {}))
        if join:
            self.worker_proc.join()
        if self.exit_hook:
            self.exit_hook()
            self.exit_hook = None

    def __del__(self):
        # noinspection PyBroadException
        try:
            self.exit(join=False)
        except Exception:
            pass
        else:
            try_run(self.worker_proc.join)


def _worker_proc_loop(
    epoch: int,
    buffer_size: int,
    dataset_dict: Dict[str, Any],
    parent_conn: mpConnection,
):
    if sys.platform == "linux":
        with open("/proc/self/comm", "w") as f:
            f.write(f"CFD worker {epoch}")

    assert isinstance(epoch, int) and isinstance(buffer_size, int)
    assert isinstance(dataset_dict, dict)
    assert isinstance(parent_conn, mpConnection)

    dataset = init_dataset(dataset_dict)

    got_init_seq_order = False
    cache: List[DatasetSeq] = []
    next_seq_idx = 0

    # noinspection PyShadowingNames
    def _add_to_cache():
        nonlocal next_seq_idx
        if len(cache) >= buffer_size:
            return False
        if not dataset.is_less_than_num_seqs(next_seq_idx):
            return False
        dataset.load_seqs(next_seq_idx, next_seq_idx + 1)
        seq_tag = dataset.get_tag(next_seq_idx)
        features = {data_key: dataset.get_data(next_seq_idx, data_key) for data_key in dataset.get_data_keys()}
        res = DatasetSeq(seq_idx=next_seq_idx, seq_tag=seq_tag, features=features)
        cache.append(res)
        next_seq_idx += 1
        return True

    # noinspection PyShadowingNames
    def _get_from_cache(seq_idx: int) -> Optional[DatasetSeq]:
        if not cache:
            return None
        if seq_idx > cache[-1].seq_idx:
            return None
        for seq in cache:
            assert seq.seq_idx <= seq_idx
            if seq.seq_idx == seq_idx:
                return seq
        assert False

    # noinspection PyShadowingNames
    def _get(seq_idx: int) -> Optional[DatasetSeq]:
        if cache and seq_idx < cache[0].seq_idx:
            raise Exception(
                f"requested seq idx {seq_idx} is smaller than cache start {cache[0].seq_idx}, cannot go backwards"
            )
        res = _get_from_cache(seq_idx)
        if res:
            return res
        if not dataset.is_less_than_num_seqs(seq_idx):
            return None
        assert next_seq_idx <= seq_idx
        while True:
            if not _add_to_cache():
                raise Exception(
                    f"buffer too small, requested seq idx {seq_idx},"
                    f" cache starts at {cache[0].seq_idx if cache else None}"
                )
            assert cache[-1].seq_idx <= seq_idx
            if cache[-1].seq_idx == seq_idx:
                return cache[-1]

    try:
        while True:
            if got_init_seq_order:
                while not parent_conn.poll():
                    if not _add_to_cache():
                        break
            msg, kwargs = parent_conn.recv()
            if msg == "exit":
                break
            elif msg == "get_data_seq":
                seq_idx = kwargs["seq_idx"]
                while cache and cache[0].seq_idx < seq_idx:
                    cache.pop(0)
                res = _get(seq_idx)
                parent_conn.send(("data_seq", res))
            elif msg == "init_seq_order":
                # We are responsible to get the seq order and distrib it to all the other workers.
                dataset.init_seq_order(**kwargs)
                parent_conn.send(("num_seqs", dataset.num_seqs))
                got_init_seq_order = True
                next_seq_idx = 0
                cache[:] = []
            else:
                raise Exception(f"unknown msg {msg!r}")
    except KeyboardInterrupt:  # when parent dies
        pass


class _FileCacheProc:
    def __init__(self, *, name: str, file_cache_opts: Dict[str, Any]):
        self.file_cache_opts = file_cache_opts

        parent_conn, child_conn = _mp.Pipe()
        self.parent_conn: mpConnection = parent_conn

        self.proc = _mp.Process(
            name=f"{name} file cache",
            target=_file_cache_proc_loop,
            args=(file_cache_opts, child_conn),
            daemon=True,
        )
        self.proc.start()
        # Make sure the child connection is closed here.
        # It stays open in the child, until the child dies.
        # When that happens, now any consecutive read on the pipe
        # should yield an exception -- which is what we want,
        # otherwise it would just hang.
        child_conn.close()

    def handle_cached_files_in_config(self, config: Dict[str, Any]) -> Tuple[Dict[str, Any], _FileCacheExitHook]:
        """:func:`FileCache.handle_cached_files_in_config`"""
        self.parent_conn.send(("handle_cached_files_in_config", {"config": config}))
        msg, (res, files) = self.parent_conn.recv()
        assert msg == "config_and_cached_files" and isinstance(config, dict) and isinstance(files, list)
        return res, _FileCacheExitHook(self, files)

    def release_files(self, filenames: Collection[str]):
        """:func:`FileCache.release_files`"""
        self.parent_conn.send(("release_files", {"filenames": filenames}))

    def exit(self, *, join: bool = True):
        """exit"""
        if self.proc.exitcode is not None:  # already stopped?
            return
        self.parent_conn.send(("exit", {}))
        if join:
            self.proc.join()

    def __del__(self):
        # noinspection PyBroadException
        try:
            self.exit(join=False)
        except Exception:
            pass
        else:
            try_run(self.proc.join)


class _FileCacheExitHook:
    def __init__(self, file_cache_proc: _FileCacheProc, cached_files: List[str]):
        self.file_cache_proc = file_cache_proc
        self.cached_files = cached_files

    def __call__(self):
        if self.file_cache_proc.proc.exitcode is not None:
            return
        self.file_cache_proc.release_files(self.cached_files)


def _file_cache_proc_loop(
    file_cache_opts: Dict[str, Any],
    parent_conn: mpConnection,
):
    if sys.platform == "linux":
        with open("/proc/self/comm", "w") as f:
            f.write(f"File cache")

    assert isinstance(file_cache_opts, dict)
    assert isinstance(parent_conn, mpConnection)

    from returnn.util.file_cache import FileCache

    file_cache = FileCache(**file_cache_opts)

    try:
        while True:
            msg, kwargs = parent_conn.recv()
            if msg == "exit":
                break
            elif msg == "handle_cached_files_in_config":
                res, files = file_cache.handle_cached_files_in_config(**kwargs)
                parent_conn.send(("config_and_cached_files", (res, files)))
            elif msg == "release_files":
                file_cache.release_files(**kwargs)
            else:
                raise Exception(f"unknown msg {msg!r}")
    except KeyboardInterrupt:  # when parent dies
        pass
