from __future__ import annotations
from typing import List
import gc
import sys
import time
import numpy
import threading
import typing
from .basic import Dataset
from returnn.log import log
from returnn.util import NumbersDict


class CachedDataset(Dataset):
    """
    Base class for datasets with caching. This is only used for the :class:`HDFDataset`.
    Also see :class:`CachedDataset2`.
    """

    def __init__(self, cache_byte_size=0, **kwargs):
        """
        :param int cache_byte_size:
        """
        super(CachedDataset, self).__init__(**kwargs)
        self._cache_byte_size = cache_byte_size
        self.cache_byte_size_total_limit = cache_byte_size
        if cache_byte_size == -1:
            self.cache_byte_size_limit_at_start = 1024**4
        elif cache_byte_size == 0:
            self.cache_byte_size_limit_at_start = 0
        else:
            self.cache_byte_size_limit_at_start = max(cache_byte_size * 2 // 3, 1)
            self.cache_byte_size_total_limit = max(cache_byte_size - self.cache_byte_size_limit_at_start, 1)
        self.num_seqs_cached_at_start = 0
        self.cached_bytes_at_start = 0
        self.nbytes = 0
        self.start_cache_initialized = False
        self.definite_cache_leftover = 0
        self.cache_num_frames_free = 0
        self.preload_set = set([])
        self.preload_end = 0
        self.alloc_intervals = None  # type: typing.Optional[list]
        self._seq_start = []  # [numpy.array([0,0])]  # uses sorted seq idx, see set_batching()
        self._seq_index = []  # type: typing.List[int]  # Via init_seq_order(). seq_index idx -> hdf seq idx
        self._seq_index_inv = {}  # type: typing.Dict[int,int]  # Via init_seq_order(). hdf seq idx -> seq_index idx
        self._index_map = range(len(self._seq_index))  # sorted seq idx -> seq_index idx
        self._tag_idx = {}  # type: typing.Dict[str,int]  # map of tag -> real-seq-idx. call _update_tag_idx
        self.targets = {}
        self.target_keys = (
            []
        )  # the keys for which we provide data; we may have labels for additional keys in self.labels
        self.timestamps = None

    def initialize(self):
        """
        Initialization.
        """
        super(CachedDataset, self).initialize()

        if self.cache_byte_size_limit_at_start > 0:
            self.nbytes = numpy.array([], dtype=numpy.float32).itemsize * (self.num_inputs * self.window + 1 + 1)

            # Calculate cache sizes.
            temp_cache_size_bytes = max(0, self.cache_byte_size_total_limit)
            self.definite_cache_leftover = (
                temp_cache_size_bytes if self.num_seqs_cached_at_start == self.num_seqs else 0
            )
            self.cache_num_frames_free = temp_cache_size_bytes // self.nbytes

            print(
                "cached %i seqs" % self.num_seqs_cached_at_start,
                "%s GB" % (self.cached_bytes_at_start / float(1024 * 1024 * 1024)),
                ("(fully loaded, %s GB left over)" if self.definite_cache_leftover else "(%s GB free)")
                % max(temp_cache_size_bytes / float(1024 * 1024 * 1024), 0),
                file=log.v4,
            )

    def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
        """
        :type epoch: int|None
        :param list[str]|None seq_list: List of sequence tags, to set a predefined order.
        :param list[int]|None seq_order: List of corpus sequence indices, to set a predefined order.
        Initialize lists:
          self.seq_index  # sorted seq idx
        """
        super(CachedDataset, self).init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
        if seq_order is not None:
            seq_index = seq_order
        elif seq_list is not None:
            self._update_tag_idx()
            seq_index = [self._tag_idx[tag] for tag in seq_list]
        else:
            seq_index = self.get_seq_order_for_epoch(
                epoch, self._num_seqs, lambda s: self._get_seq_length_by_real_idx(s)[0]
            )

        old_index_map = self._index_map[:]
        self._index_map = range(len(seq_index))  # sorted seq idx -> seq_index idx

        if isinstance(seq_index, numpy.ndarray) or isinstance(self._seq_index, numpy.ndarray):
            seq_index_unchanged = numpy.array_equal(self._seq_index, seq_index)
        else:
            seq_index_unchanged = self._seq_index == seq_index

        if seq_index_unchanged and self.start_cache_initialized:
            return False

        if epoch is not None:
            # Give some hint to the user in case he is wondering why the cache is reloading.
            print("Reinitialize dataset seq order for epoch %i." % epoch, file=log.v4)

        if (
            self.cache_byte_size_limit_at_start == 0
            or self.num_seqs_cached_at_start != len(seq_index)
            or not self.start_cache_initialized
        ):
            self._seq_index = seq_index
            self._seq_index_inv = {}  # reset, create later if needed
            self._init_seq_starts()
            self._init_alloc_intervals()
            self._init_start_cache()
            self.start_cache_initialized = True
        else:
            if not self._seq_index_inv:
                self._seq_index_inv = dict(
                    zip(self._seq_index, range(len(self._seq_index)))
                )  # hdf seq idx -> seq_index idx
            self._index_map = [self._seq_index_inv[i] for i in seq_index]  # sorted seq idx -> seq_index idx
            if self._index_map == old_index_map:
                return False
        return True

    def supports_seq_order_sorting(self) -> bool:
        """supports sorting"""
        return True

    def supports_sharding(self) -> bool:
        """supports sharding"""
        return True

    def get_current_seq_order(self):
        assert self.cache_byte_size_limit_at_start == 0  # not implemented otherwise, we ignore _index_map
        return self._seq_index

    def _get_tag_by_real_idx(self, real_idx):
        raise NotImplementedError

    def _update_tag_idx(self):
        if self._tag_idx:
            return
        for i in range(self._num_seqs):
            self._tag_idx[self._get_tag_by_real_idx(i)] = i

    def batch_set_generator_cache_whole_epoch(self):
        return True

    def _init_alloc_intervals(self):
        if self.cache_byte_size_limit_at_start == 0:
            return
        if self.epoch is None:
            return
        assert self.num_seqs > 0
        assert self.num_inputs > 0
        assert self.window > 0
        self.preload_set = set([])
        self.alloc_intervals = [
            (0, 0, numpy.zeros([1] + self.get_data_shape("data"), dtype=self.get_data_dtype("data"))),
            (
                self.num_seqs,
                self.num_seqs,
                numpy.zeros([1] + self.get_data_shape("data"), dtype=self.get_data_dtype("data")),
            ),
        ]
        # self.alloc_intervals[i] is (idx start, idx end, data), where
        # idx start/end is the sorted seq idx start/end, end exclusive,
        # and data is a numpy.array.

    def _init_seq_starts(self):
        if self.cache_byte_size_limit_at_start == 0:
            return
        self._seq_start = [self._seq_start[0] * 0]  # idx like in seq_index, *not* real idx
        for i in range(self.num_seqs):
            ids = self._seq_index[i]
            self._seq_start.append(self._seq_start[-1] + self._get_seq_length_by_real_idx(ids))

    def _init_start_cache(self):
        if self.cache_byte_size_limit_at_start == 0:
            return
        if not self.alloc_intervals:
            return
        if not self.nbytes:
            return
        if not self.epoch:
            return

        num_cached = 0
        cached_bytes = 0
        for i in range(self.num_seqs):
            if i == num_cached:
                nbytes = self.get_seq_length_nd(i)[0] * self.nbytes
                if self.cache_byte_size_limit_at_start >= cached_bytes + nbytes:
                    num_cached = i + 1
                    cached_bytes += nbytes

        self.num_seqs_cached_at_start = num_cached
        self.cached_bytes_at_start = cached_bytes
        if num_cached > 0:
            self.preload_end = num_cached
            if sys.version_info >= (3, 0):
                threading.Thread(target=self._preload_seqs, args=(0, num_cached), daemon=True).start()
            else:
                threading.Thread(target=self._preload_seqs, args=(0, num_cached)).start()

    def load_seqs(self, start, end):
        """
        Load data sequences.
        As a side effect, will modify / fill-up:
          self.alloc_intervals
          self.targets
        This does some extra logic for the cache and calls self._load_seqs()
        for the real loading.

        :param int start: start sorted seq idx
        :param int end: end sorted seq idx
        """
        assert start >= 0
        assert start <= end

        if self.is_cached(start, end, blocking=False):
            return

        if self.cache_byte_size_limit_at_start > 0:  # If the cache is enabled.
            self._load_seqs_with_cache(start, end)
            return self.is_cached(start, end, blocking=True)

        super(CachedDataset, self).load_seqs(start, end)

    def _load_seqs(self, start, end):
        raise NotImplementedError

    def _load_seqs_with_cache(self, start, end, clear=True):
        if not clear:
            # only remove as many frames as required
            num_needed_cache_frames = self.get_seq_start(end)[0] - self.get_seq_start(start)[0]
            if self.cache_num_frames_free < num_needed_cache_frames:
                self.cache_num_frames_free += self.delete(num_needed_cache_frames - self.cache_num_frames_free)
                gc.collect()
            self.cache_num_frames_free -= num_needed_cache_frames
            threading.Thread(target=self._preload_seqs, args=(start, end)).start()
        else:
            # First, delete everything.
            self.cache_num_frames_free += self.delete(None)
            gc.collect()
            # Preload as much as we can so that we fill up the cache.
            while end < self.num_seqs:
                num_needed_cache_frames = self.get_seq_length_nd(end)[0]
                if self.cache_num_frames_free - num_needed_cache_frames < 0:
                    break
                self.cache_num_frames_free -= num_needed_cache_frames
                end += 1
            self.preload_end = end
            threading.Thread(target=self._preload_seqs, args=(start, end)).start()

    def _preload_seqs(self, start, end):
        print("Preloading cache from", start, "to", end, file=log.v4)
        super(CachedDataset, self).load_seqs(start, end)
        self.preload_end = self.num_seqs_cached_at_start

    def _shuffle_frames_in_seqs(self, start, end):
        """
        :type start: int
        :type end: int
        """
        assert start < end
        assert self.is_cached(start, end)
        alloc_idx = self.alloc_interval_index(start)
        alloc_start, alloc_end, alloc_data = self.alloc_intervals[alloc_idx]
        assert start >= alloc_start
        assert end <= alloc_end
        rnd = numpy.random.RandomState(start)  # Some deterministic way to shuffle!
        num_frames = self._seq_start[end][0] - self._seq_start[start][0]
        assert num_frames > 0
        perm = rnd.permutation(num_frames)
        alloc_offset = self._seq_start[start][0] - self._seq_start[alloc_start][0]
        assert alloc_offset + num_frames <= alloc_data.shape[0]
        # Permute alloc_data.
        data = alloc_data[alloc_offset : alloc_offset + num_frames]
        alloc_data[alloc_offset : alloc_offset + num_frames] = data[perm]
        # Permute targets.
        for k in self.targets:
            idx = self.target_keys.index(k) + 1
            targets = self.targets[k][self._seq_start[idx] : self._seq_start[start][idx] + num_frames]
            self.targets[k][
                self._seq_start[start][idx] : self._seq_start[start][idx]
                + self._seq_start[end][idx]
                - self._seq_start[start][idx]
            ] = targets[perm]

    def _set_alloc_intervals_data(self, idc, data):
        """
        :param int idc: index of sorted seq idx
        :param numpy.ndarray data: raw data
        """
        idi = self.alloc_interval_index(idc)
        assert idi >= 0
        o = self._seq_start[idc][0] - self._seq_start[self.alloc_intervals[idi][0]][0]
        l = data.shape[0]
        x = data
        if self.window > 1:
            x = self._sliding_window(x)
        self.alloc_intervals[idi][2][o : o + l] = x

    def alloc_interval_index(self, ids):
        """
        :param int ids: sorted seq idx
        :return index in self.alloc_intervals
        :rtype: int
        """
        s = 0
        e = len(self.alloc_intervals)
        # Binary search.
        while s < e:
            i = (s + e) // 2
            alloc_start, alloc_end, _ = self.alloc_intervals[i]
            if alloc_start <= ids < alloc_end:
                return i
            elif alloc_start <= ids and ids >= alloc_end:
                if s == i:
                    return -1
                s = i
            elif alloc_start > ids:
                if e == i:
                    return -1
                e = i
            else:
                assert False
        return -1

    def _insert_alloc_interval(self, pos, value, merge=False):
        """
        Insert np.zeros into self.alloc_intervals.
        :param int pos: idx in self.alloc_intervals
        :param (int,int) value: (start,end) like in load_seqs(), sorted seq idx
        :rtype: int
        """
        if value[0] == value[1]:
            return 0
        ci = self.alloc_intervals[pos][1]
        ni = self.alloc_intervals[pos + 1][0]
        xc = self.alloc_intervals[pos][2]
        xn = self.alloc_intervals[pos + 1][2]
        if value[0] == ci and value[1] == ni and merge:
            nj = self.alloc_intervals[pos][0]
            nk = self.alloc_intervals[pos + 1][1]
            del self.alloc_intervals[pos]
            del self.alloc_intervals[pos]
            self.alloc_intervals.insert(
                pos,
                (
                    nj,
                    nk,
                    numpy.concatenate(
                        [
                            xc,
                            numpy.zeros(
                                [self._seq_start[ni][0]] + self.get_data_shape("data"),
                                dtype=self.get_data_dtype("data"),
                            ),
                            xn,
                        ]
                    ),
                ),
            )
            return 0
        elif value[0] == ci and merge:
            nj = self.alloc_intervals[pos][0]
            del self.alloc_intervals[pos]
            self.alloc_intervals.insert(
                pos,
                (
                    nj,
                    value[1],
                    numpy.concatenate(
                        [
                            xc,
                            numpy.zeros(
                                [self._seq_start[value[1]][0] - self._seq_start[ci][0]] + self.get_data_shape("data"),
                                dtype=self.get_data_dtype("data"),
                            ),
                        ]
                    ),
                ),
            )
            return 0
        elif value[1] == ni and merge:
            nk = self.alloc_intervals[pos + 1][1]
            del self.alloc_intervals[pos + 1]
            self.alloc_intervals.insert(
                pos + 1,
                (
                    value[0],
                    nk,
                    numpy.concatenate(
                        [
                            numpy.zeros(
                                [self._seq_start[ni][0] - self._seq_start[value[0]][0]] + self.get_data_shape("data"),
                                dtype=self.get_data_dtype("data"),
                            ),
                            xc,
                        ]
                    ),
                ),
            )
            return 0
        else:
            self.alloc_intervals.insert(
                pos + 1,
                value
                + (
                    numpy.zeros(
                        [self._seq_start[value[1]][0] - self._seq_start[value[0]][0]] + self.get_data_shape("data"),
                        dtype=self.get_data_dtype("data"),
                    ),
                ),
            )
            return 1

    def _remove_alloc_interval(self, pos, value):
        """
        Remove data from self.alloc_intervals.
        :param int pos: idx in self.alloc_intervals
        :param (int,int) value: (start,end) like in load_seqs(), sorted seq idx
        :rtype: int
        """
        ci, ni, xi = self.alloc_intervals[pos]
        if value[0] == ci and value[1] == ni:
            del self.alloc_intervals[pos]
            return -1
        elif value[0] == ci:
            self.alloc_intervals.insert(
                pos, (value[1], ni, xi[self._seq_start[value[1]][0] - self._seq_start[ci][0] :])
            )
            del self.alloc_intervals[pos + 1]
            return 0
        elif value[1] == ni:
            self.alloc_intervals.insert(
                pos, (ci, value[0], xi[: self._seq_start[value[0]][0] - self._seq_start[ci][0]])
            )
            del self.alloc_intervals[pos + 1]
            return 0
        else:
            self.alloc_intervals.insert(
                pos, (value[1], ni, xi[self._seq_start[value[1]][0] - self._seq_start[ci][0] :])
            )
            self.alloc_intervals.insert(
                pos, (ci, value[0], xi[: self._seq_start[value[0]][0] - self._seq_start[ci][0]])
            )
            del self.alloc_intervals[pos + 2]
            return 1

    def _modify_alloc_intervals(self, start, end, invert):
        """
        Inserts/removes sorted seq idx range (start,end).
        :param int start: like in load_seqs(), sorted seq idx
        :param int end: like in load_seqs(), sorted seq idx
        :param bool invert: True->insert, False->remove
        :rtype: list[int]
        :return selection list, modified sorted seq idx in self.alloc_intervals
        """
        if end is None:
            end = start + 1
        if start == end:
            return
        assert start < end
        i = 0
        selection = []  # type: typing.List[int]
        modify = self._insert_alloc_interval if invert else self._remove_alloc_interval
        while i < len(self.alloc_intervals) - invert:
            ni = self.alloc_intervals[i + invert][1 - invert]  # insert mode: start idx of next alloc
            ci = self.alloc_intervals[i][invert]  # insert mode: end idx of cur alloc
            assert ci <= ni
            flag = ((ci <= start < ni), (ci < end <= ni), (ci < start and ni <= start) or (ci >= end and ni > end))
            if not flag[0] and not flag[1]:
                if not flag[2]:
                    selection.extend(range(ci, ni))
                    i += modify(i, (ci, ni))
            elif flag[1]:
                v = (start if flag[0] else ci, end)
                selection.extend(range(v[0], v[1]))
                i += modify(i, v)
                break
            elif flag[0]:
                selection.extend(range(start, ni))
                i += modify(i, (start, ni))
            i += 1
        if self.alloc_intervals[0][0] != 0:
            self.alloc_intervals.insert(
                0, (0, 0, numpy.zeros([1] + self.get_data_shape("data"), dtype=self.get_data_dtype("data")))
            )
        if self.alloc_intervals[-1][1] != self.num_seqs:
            self.alloc_intervals.append(
                (
                    self.num_seqs,
                    self.num_seqs,
                    numpy.zeros([1] + self.get_data_shape("data"), dtype=self.get_data_dtype("data")),
                )
            )
        return selection

    def insert_alloc_interval(self, start, end=None):
        return self._modify_alloc_intervals(start, end, True)

    def remove_alloc_interval(self, start, end=None):
        return self._modify_alloc_intervals(start, end, False)

    def delete(self, nframes):
        """
        :param int|None nframes: how much frames to delete max.
          Note that this limit is not strict. We can end up
          deleting more than nframes.
        :return: number of frames deleted
        :rtype: int
        """
        if nframes is not None:
            if nframes == 0:
                return 0
            assert nframes > 0
        deleted = 0
        i = 0
        while (not nframes or deleted < nframes) and i < len(self.alloc_intervals):
            ai = self.alloc_intervals[i]
            if ai[1] > self.num_seqs_cached_at_start and ai[0] < ai[1]:
                removed = self.remove_alloc_interval(max(ai[0], self.num_seqs_cached_at_start), ai[1])
                self.preload_set -= set(removed)
                deleted += sum([self._get_seq_length_by_real_idx(self._seq_index[i])[0] for i in removed])
            else:
                i += 1
        return deleted

    @property
    def num_seqs(self):
        return len(self._index_map)

    def is_cached(self, start, end, blocking=False):
        """
        :param int start: like in load_seqs(), sorted seq idx
        :param int end: like in load_seqs(), sorted seq idx
        :rtype: bool
        :returns whether we have the full range (start,end) of sorted seq idx
          cached in self.alloc_intervals (end is exclusive).
        """
        if self.cache_byte_size_total_limit == 0:  # disabled cache
            return False
        if start == end:
            return True  # Empty.
        assert start < end
        if blocking and end <= self.preload_end:
            while not set(range(start, end)) <= self.preload_set:
                time.sleep(0.2)
            return True
        return set(range(start, end)) <= self.preload_set

    def _get_seq_length_by_real_idx(self, real_seq_idx):
        """
        :param int real_seq_idx:
        :returns length of the sequence with index 'real_seq_idx'
        :rtype: numpy.ndarray
        """
        raise NotImplementedError

    def get_seq_length_nd(self, sorted_seq_idx):
        """
        :type sorted_seq_idx: int
        :rtype: numpy.ndarray
        """
        real_seq_idx = self._seq_index[self._index_map[sorted_seq_idx]]
        return self._get_seq_length_by_real_idx(real_seq_idx)

    def get_seq_length(self, seq_idx):
        """
        :rtype: NumbersDict
        """
        lengths = self.get_seq_length_nd(seq_idx)
        d = {}
        first_target_idx = 0
        # We allow using only targets. In this case self.num_inputs == 0 and the "data" key is not used.
        if self.num_inputs > 0:
            d["data"] = lengths[0]
            first_target_idx = 1
        for k, l in zip(self.target_keys, lengths[first_target_idx:]):
            d[k] = l
        return NumbersDict(d)

    def get_seq_start(self, sorted_seq_idx):
        """
        :type sorted_seq_idx: int
        :rtype: (int,int)
        """
        return self._seq_start[sorted_seq_idx]

    def get_times(self, sorted_seq_idx):
        seq_start = self.get_seq_start(sorted_seq_idx)[0]
        seq_len = self.get_seq_length_nd(sorted_seq_idx)[0]
        return self.timestamps[seq_start : seq_start + seq_len]

    def get_input_data(self, sorted_seq_idx):
        seq_idx = self._index_map[sorted_seq_idx]
        idi = self.alloc_interval_index(seq_idx)
        assert idi >= 0, "failed to get data for seq %i" % sorted_seq_idx
        alloc_start_seq, alloc_end_seq, alloc_data = self.alloc_intervals[idi]
        o = self.get_seq_start(seq_idx)[0] - self.get_seq_start(alloc_start_seq)[0]
        assert o >= 0
        l = self.get_seq_length_nd(sorted_seq_idx)[0]
        assert alloc_data.shape[0] >= o + l
        return alloc_data[o : o + l]

    def get_data_dim(self, key):
        if key == "data" and self.num_inputs > 0:  # if num_inputs == 0, we allow "data" as a target key
            return self.num_inputs * self.window
        return self.num_outputs[key][0]

    def get_targets(self, target, sorted_seq_idx):
        seq_idx = self._index_map[sorted_seq_idx]
        idx = self.target_keys.index(target) + 1
        seq_start = self.get_seq_start(seq_idx)[idx]
        seq_len = self.get_seq_length_nd(sorted_seq_idx)[idx]
        return self.targets[target][seq_start : seq_start + seq_len]

    def get_data_keys(self) -> List[str]:
        """data keys"""
        res = []
        if self.num_inputs > 0:
            res.append("data")
        res.extend(self.get_target_list())
        return res

    def get_target_list(self):
        return list(self.targets.keys())

    def get_tag(self, sorted_seq_idx):
        raise NotImplementedError

    def have_corpus_seq_idx(self):
        return True

    def get_corpus_seq_idx(self, seq_idx):
        """
        :param int seq_idx: sorted sequence index from the current epoch, depending on seq_ordering
        :return: the sequence index as-is in the original corpus. only defined if self.have_corpus_seq_idx()
        :rtype: int
        """
        return self._seq_index[self._index_map[seq_idx]]
