# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Load nnet3 training egs which generated by kaldi
"""

import random
import numpy as np
import kaldi_python_io as k_io
from paddle.io import Dataset
import paddlespeech.vector.utils.utils as utils
from paddlespeech.vector import _logger as log
class KaldiEgsDataset(Dataset):
    """
    Dataset used to load kaldi nnet3 egs files.
    """
    def __init__(self, egs_list_file, egs_idx, transforms=None):
        self.scp_reader = None
        self.subset_idx = egs_idx - 1
        self.transforms = transforms
        if not utils.is_exist(egs_list_file):
            return

        self.egs_files = []
        with open(egs_list_file, 'r') as in_fh:
            for line in in_fh:
                if line.strip():
                    self.egs_files.append(line.strip())

        self.next_subset()

    def next_subset(self, target_index=None, delta_index=None):
        """
        Use next specific subset

        Args:
            target_index: target egs index
            delta_index: incremental value of egs index
        """
        if self.egs_files:
            if target_index:
                self.subset_idx = target_index
            else:
                delta_index = delta_index if delta_index else 1
                self.subset_idx += delta_index
            log.info("egs dataset subset index: %d" % (self.subset_idx))
            egs_file = self.egs_files[self.subset_idx % len(self.egs_files)]
            if utils.is_exist(egs_file):
                self.scp_reader = k_io.Nnet3EgsScriptReader(egs_file)
            else:
                log.warning("No such file or directory: %s" % (egs_file))

    def __getitem__(self, index):
        if self.scp_reader is None:
            return {}
        index %= len(self)
        in_dict, out_dict = self.scp_reader[index]
        x = np.array(in_dict['matrix'])
        x = np.transpose(x)
        y = np.array(out_dict['matrix'][0][0][0], dtype=np.int).reshape((1,))
        if self.transforms is not None:
            idx = random.randint(0, len(self.transforms) - 1)
            x = self.transforms[idx](x)
        return x, y

    def __len__(self):
        return len(self.scp_reader)

    def __iter__(self):
        self._start = 0
        return self

    def __next__(self):
        if self._start < len(self):
            ret = self[self._start]
            self._start += 1
            return ret
        else:
            raise StopIteration