#   Copyright (c) 2020 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.
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import io
import os
import sys
import six
import random

import paddle
import paddle.fluid as fluid
import numpy as np


def word2id(word_dict, query):
    """
    Convert word sequence into id list
    """
    unk_id = len(word_dict)
    wids = [
        word_dict[w] if w in word_dict else unk_id
        for w in query.strip().split(" ")
    ]
    return wids


def pad_wid(wids, max_seq_len=128, pad_id=0):
    """
    Padding data to max_seq_len
    """
    seq_len = len(wids)
    if seq_len < max_seq_len:
        for i in range(max_seq_len - seq_len):
            wids.append(pad_id)
    else:
        wids = wids[:max_seq_len]
    return wids


def data_reader(file_path, word_dict, num_examples, phase, epoch, max_seq_len, shuffle=False):
    """
    Data reader, which convert word sequence into id list
    """
    unk_id = len(word_dict)
    all_data = []
    with io.open(file_path, "r", encoding='utf8') as fin:
        for line in fin:
            if line.startswith("label"):
                continue
            if phase == "infer":
                cols = line.strip().split("\t")
                query = cols[-1] if len(cols) != -1 else cols[0]
                wids = word2id(word_dict, query)
                wids = pad_wid(wids, max_seq_len, unk_id)
                all_data.append((wids))
            else:
                cols = line.strip().split("\t")
                if len(cols) != 2:
                    sys.stderr.write("[NOTICE] Error Format Line!")
                    continue
                label = [int(cols[0])]
                query = cols[1].strip()
                wids = word2id(word_dict, query)
                wids = pad_wid(wids, max_seq_len, unk_id)
                all_data.append((wids, label))
    num_examples[phase] = len(all_data)

    if phase == "infer":

        def reader():
            """
            Infer reader function
            """
            for wids in all_data:
                yield wids

        return reader

    def reader():
        """
        Reader function
        """
        for idx in range(epoch):
            if phase == "train" and shuffle:
                random.shuffle(all_data)
            for wids, label in all_data:
                yield wids, label

    return reader


def load_vocab(file_path):
    """
    load the given vocabulary
    """
    vocab = {}
    with io.open(file_path, 'r', encoding='utf8') as fin:
        wid = 0
        for line in fin:
            if line.strip() not in vocab:
                vocab[line.strip()] = wid
                wid += 1
    vocab["<unk>"] = len(vocab)
    return vocab

def query2ids(vocab_path, query):
    """
    Convert query to id list according to the given vocab
    """
    vocab = load_vocab(vocab_path)
    wids = word2id(vocab, query)
    return wids
