"""
Uses KenLM (https://kheafield.com/code/kenlm/) (extern/kenlm) to read n-gram LMs (ARPA format),
and provides a TF op to use them.

"""

import sys
import os
import tensorflow as tf

returnn_dir = os.path.dirname(os.path.abspath(__file__))
kenlm_dir = returnn_dir + "/extern/kenlm"


def kenlm_checked_out():
    """
    :rtype: bool
    """
    return os.path.exists("%s/lm/test.arpa" % kenlm_dir)


# https://www.tensorflow.org/guide/extend/op
# Also see TFUtil.TFArrayContainer for TF resources.
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/tensor.h
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/tensor_shape.h
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/tensor_types.h
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/lib/strings/str_util.h
_src_code = """
#include <exception>
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/resource_op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/public/version.h"


using namespace tensorflow;

#if (TF_MAJOR_VERSION < 2) || (TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION < 2)
using tstring = std::string;
#else
using tstring = tensorflow::tstring;
#endif


REGISTER_OP("KenLmLoadModel")
.Attr("filename: string")
.Attr("container: string = ''")
.Attr("shared_name: string = ''")
.Output("handle: resource")
.SetIsStateful()
.SetShapeFn(shape_inference::ScalarShape)
.Doc("KenLmLoadModel: loads KenLM model, creates TF resource, persistent across runs in the session");


REGISTER_OP("KenLmAbsScoreStrings")
.Input("handle: resource")
.Input("strings: string")
.Output("scores: float32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
  c->set_output(0, c->input(1));
  return Status::OK();
})
.Doc("KenLmScoreStrings: scores texts. returns in +log space (natural log, not base 10)");


REGISTER_OP("KenLmAbsScoreBpeStrings")
.Input("handle: resource")
.Input("bpe_merge_symbol: string")
.Input("strings: string")
.Output("scores: float32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
  c->set_output(0, c->input(2));
  return Status::OK();
})
.Doc("KenLmAbsScoreBpeStrings: optionally BPE-merges, remove surrounding whitespaces and scores texts."
  " returns in +log space (natural log, not base 10)."
  " relative score, relative to previous text."
  );


REGISTER_OP("KenLmAbsScoreBpeStringsDense")
.Input("handle: resource")
.Input("bpe_merge_symbol: string")
.Input("strings: string")
.Input("labels: string")
.Output("scores: float32")
.Output("dense_scores: float32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
  c->set_output(0, c->input(2));
  ::tensorflow::shape_inference::ShapeHandle out_shape;
  TF_RETURN_IF_ERROR(c->Concatenate(c->input(2), c->input(3), &out_shape));
  c->set_output(1, out_shape);
  return Status::OK();
})
.Doc("KenLmAbsScoreBpeStrings: optionally BPE-merges, remove surrounding whitespaces and scores texts."
  " returns in +log space (natural log, not base 10)."
  " relative score, relative to previous text."
  " dense output, for all possible succeeding labels.");


// https://github.com/kpu/kenlm/blob/master/lm/model.hh
// https://github.com/kpu/kenlm/blob/master/lm/virtual_interface.hh
// https://github.com/kpu/kenlm/blob/master/python/kenlm.pyx
struct KenLmModel : public ResourceBase {
  explicit KenLmModel(const string& filename)
      : filename_(filename), model_(filename.c_str()) {}

  float abs_score(const ::tstring& text) {
    float total = 0;
    mutex_lock l(mu_);
    lm::ngram::State state, out_state;
    model_.BeginSentenceWrite(&state);
    for(const string& word : tensorflow::str_util::Split(text, ' ')) {
      if(word.empty()) continue;
      auto word_idx = model_.BaseVocabulary().Index(word);
      total += model_.FullScore(state, word_idx, out_state).prob;
      state = out_state;
    }
    // KenLM returns score in +log10 space.
    // We want to return in (natural) +log space.
    // 10 ** x = e ** (x * log(10))
    return total * logf(10.);
  }

  // See comments below.
  // We expect that the text either ends with a space or not, i.e. "... word " or "... subword".
  float abs_score_dense(
        const ::tstring& text, const ::tstring& last_word_join,
        const TTypes<::tstring>::ConstFlat labels, TTypes<float>::UnalignedFlat out_dense_scores) {
    assert(labels.size() == out_dense_scores.size());
    mutex_lock l(mu_);
    lm::ngram::State state, out_state;
    model_.BeginSentenceWrite(&state);
    // We expect that the text either ends with a space or not, i.e. "... word " or "... subword".
    // We split the text into words. In the first case, we would have an empty word at the end, otherwise not.
    auto words = tensorflow::str_util::Split(text, ' ');
    float total_score = 0;
    ::tstring last_word = "";
    if(!words.empty()) {
      last_word = words[words.size() - 1];
      // Only up to the last word, which is either empty or a subword, which we join below.
      for(int i = 0; i < words.size() - 1; ++i) {
        auto word = words[i];
        if(word.empty()) continue;
        auto word_idx = model_.BaseVocabulary().Index(word);
        total_score += model_.FullScore(state, word_idx, out_state).prob;
        state = out_state;
      }
    }
    for(int i = 0; i < labels.size(); ++i) {
      ::tstring word = last_word + labels(i);
      auto word_idx = model_.BaseVocabulary().Index(word);
      float score = model_.FullScore(state, word_idx, out_state).prob;
      out_dense_scores(i) = (total_score + score) * logf(10.);
    }
    // Return the score from the prev step.
    if(!last_word.empty()) {
      auto word_idx = model_.BaseVocabulary().Index(last_word + last_word_join);
      total_score += model_.FullScore(state, word_idx, out_state).prob;
    }
    return total_score * logf(10.);
  }

  string DebugString()
#if (TF_MAJOR_VERSION == 1 && TF_MINOR_VERSION >= 14) || (TF_MAJOR_VERSION > 1)
const
#endif
  override {
    return strings::StrCat("KenLmModel[", filename_, "]");
  }

  const string filename_;
  mutex mu_;
  lm::ngram::ProbingModel model_;
};


// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/resource_op_kernel.h
// TFUtil.TFArrayContainer
class KenLmLoadModelOp : public ResourceOpKernel<KenLmModel> {
 public:
  explicit KenLmLoadModelOp(OpKernelConstruction* context)
      : ResourceOpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("filename", &filename_));
  }

 private:
  virtual bool IsCancellable() const { return false; }
  virtual void Cancel() {}

  Status CreateResource(KenLmModel** ret) override {
    try {
      *ret = new KenLmModel(filename_);
    } catch (std::exception& exc) {
      return errors::Internal("Could not load KenLmModel ", filename_, ", exception: ", exc.what());
    }
    if(*ret == nullptr)
      return errors::ResourceExhausted("Failed to allocate");
    return Status::OK();
  }

  Status VerifyResource(KenLmModel* lm) override {
    if(lm->filename_ != filename_)
      return errors::InvalidArgument("Filename mismatch: expected ", filename_,
                                     " but got ", lm->filename_, ".");
    return Status::OK();
  }

  string filename_;
};

REGISTER_KERNEL_BUILDER(Name("KenLmLoadModel").Device(DEVICE_CPU), KenLmLoadModelOp);


class KenLmAbsScoreStringsOp : public OpKernel {
 public:
  using OpKernel::OpKernel;

  void Compute(OpKernelContext* context) override {
    KenLmModel* lm;
    {
      const Tensor* handle;
      OP_REQUIRES_OK(context, context->input("handle", &handle));
      OP_REQUIRES_OK(context, GetResourceFromContext(context, "handle", &lm));
    }
    core::ScopedUnref unref(lm);

    const Tensor& input_tensor = context->input(1);
    auto input_flat = input_tensor.flat<::tstring>();

    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(), &output_tensor));
    auto output_flat = output_tensor->flat<float>();

    for(int i = 0; i < input_flat.size(); ++i) {
      output_flat(i) = lm->abs_score(input_flat(i));
    }
  }
};

REGISTER_KERNEL_BUILDER(Name("KenLmAbsScoreStrings").Device(DEVICE_CPU), KenLmAbsScoreStringsOp);


class KenLmAbsScoreBpeStringsOp : public OpKernel {
 public:
  using OpKernel::OpKernel;

  void Compute(OpKernelContext* context) override {
    KenLmModel* lm;
    {
      const Tensor* handle;
      OP_REQUIRES_OK(context, context->input("handle", &handle));
      OP_REQUIRES_OK(context, GetResourceFromContext(context, "handle", &lm));
    }
    core::ScopedUnref unref(lm);

    OP_REQUIRES(context, context->input(1).NumElements() == 1,
      errors::InvalidArgument(
        "bpe_merge_symbol must be a single element but got shape ",
        context->input(1).shape().DebugString()));
    const auto& bpe_merge_symbol = context->input(1).flat<::tstring>()(0);

    const Tensor& input_tensor = context->input(2);
    auto input_flat = input_tensor.flat<::tstring>();

    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(), &output_tensor));
    auto output_flat = output_tensor->flat<float>();

    for(int i = 0; i < input_flat.size(); ++i) {
      ::tstring text = input_flat(i);
      if(!bpe_merge_symbol.empty())
        text = tensorflow::str_util::StringReplace(text, bpe_merge_symbol + " ", "", /* replace_all */ true);
      tensorflow::StringPiece sp(text);
      tensorflow::str_util::RemoveWhitespaceContext(&sp);
      text = std::string(sp.data(), sp.size());
      output_flat(i) = lm->abs_score(text);
    }
  }
};

REGISTER_KERNEL_BUILDER(Name("KenLmAbsScoreBpeStrings").Device(DEVICE_CPU), KenLmAbsScoreBpeStringsOp);


class KenLmAbsScoreBpeStringsDenseOp : public OpKernel {
 public:
  using OpKernel::OpKernel;

  void Compute(OpKernelContext* context) override {
    KenLmModel* lm;
    {
      const Tensor* handle;
      OP_REQUIRES_OK(context, context->input("handle", &handle));
      OP_REQUIRES_OK(context, GetResourceFromContext(context, "handle", &lm));
    }
    core::ScopedUnref unref(lm);

    OP_REQUIRES(context, context->input(1).NumElements() == 1,
      errors::InvalidArgument(
        "bpe_merge_symbol must be a single element but got shape ",
        context->input(1).shape().DebugString()));
    const ::tstring& bpe_merge_symbol = context->input(1).flat<::tstring>()(0);

    const Tensor& input_tensor = context->input(2);
    auto input_flat = input_tensor.flat<::tstring>();

    const Tensor& labels_tensor = context->input(3);
    auto labels_flat = labels_tensor.flat<::tstring>();

    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(), &output_tensor));
    auto output_flat = output_tensor->flat<float>();

    Tensor* output_dense_tensor = NULL;
    TensorShape output_dense_shape(input_tensor.shape());
    output_dense_shape.AppendShape(labels_tensor.shape());
    OP_REQUIRES_OK(context, context->allocate_output(1, output_dense_shape, &output_dense_tensor));
    Tensor output_dense_flat_tensor;
    OP_REQUIRES(context,
      output_dense_flat_tensor.CopyFrom(
        *output_dense_tensor,
        TensorShape({input_tensor.NumElements(), labels_tensor.NumElements()})),
      errors::Internal("CopyFrom failed"));

    for(int i = 0; i < input_flat.size(); ++i) {
      ::tstring text = input_flat(i);
      if(!bpe_merge_symbol.empty())
        text = tensorflow::str_util::StringReplace(text, bpe_merge_symbol + " ", "", /* replace_all */ true);
      output_flat(i) = lm->abs_score_dense(
        text, bpe_merge_symbol, labels_flat, output_dense_flat_tensor.Slice(i, i + 1).unaligned_flat<float>());
    }
  }
};

REGISTER_KERNEL_BUILDER(Name("KenLmAbsScoreBpeStringsDense").Device(DEVICE_CPU), KenLmAbsScoreBpeStringsDenseOp);

"""

_kenlm_src_code_workarounds = """
// ------- start with some workarounds { ------
// The KenLM code (util/integer_to_string.cc) includes this file in the wrong namespace.
// Thus include it here now.
#include <emmintrin.h>
// ------- end with workarounds } -------------
"""


_tf_mod = None


def get_tf_mod(verbose=False):
    """
    :param bool verbose:
    :return: module
    """
    global _tf_mod
    if _tf_mod:
        return _tf_mod
    import platform
    from glob import glob
    from returnn.tf.util.basic import OpCodeCompiler

    # References:
    # https://github.com/kpu/kenlm/blob/master/setup.py
    # https://github.com/kpu/kenlm/blob/master/compile_query_only.sh

    # Collect files.
    assert kenlm_checked_out(), "submodule in %r not checked out?" % kenlm_dir
    files = glob("%s/util/*.cc" % kenlm_dir)
    files += glob("%s/lm/*.cc" % kenlm_dir)
    files += glob("%s/util/double-conversion/*.cc" % kenlm_dir)
    files = [fn for fn in files if not (fn.endswith("main.cc") or fn.endswith("test.cc"))]
    assert files, "submodule in %r not checked out?" % kenlm_dir
    libs = ["z"]
    if platform.system() != "Darwin":
        libs.append("rt")

    # Put code all together in one big blob.
    src_code = ""
    src_code += _kenlm_src_code_workarounds
    for fn in files:
        f_code = open(fn).read()
        f_code = "".join([x for x in f_code if ord(x) < 128])  # enforce ASCII
        # We need to do some replacements to not clash symbol names.
        fn_short = os.path.basename(fn).replace(".", "_")
        for word in ["kConverter"]:
            f_code = f_code.replace(word, "%s_%s" % (fn_short, word))
        src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename(fn)
        # https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control
        src_code += '#line 1 "%s"\n' % os.path.basename(fn)
        src_code += f_code
        src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename(fn)
    src_code += "\n\n// ------------ our code now: ------------\n\n"
    src_code += '#line 1 "<our code>"\n'
    src_code += _src_code

    compiler = OpCodeCompiler(
        base_name="KenLM",
        code_version=1,
        code=src_code,
        include_paths=(kenlm_dir, kenlm_dir + "/util/double-conversion"),
        c_macro_defines={"NDEBUG": 1, "KENLM_MAX_ORDER": 6, "HAVE_ZLIB": 1},
        ld_flags=["-l%s" % lib for lib in libs],
        is_cpp=True,
        use_cuda_if_available=False,
        verbose=verbose,
    )
    tf_mod = compiler.load_tf_module()
    assert hasattr(tf_mod, "ken_lm_abs_score_strings"), "content of mod: %r" % (dir(tf_mod),)
    _tf_mod = tf_mod
    return tf_mod


def ken_lm_load(filename):
    """
    :param str filename:
    :return: TF resource handle
    :rtype: tf.Tensor
    """
    return get_tf_mod().ken_lm_load_model(filename=filename)


def ken_lm_abs_score_strings(handle, strings):
    """
    :param tf.Tensor handle: TF resource handle returned by :func:`ken_lm_load`
    :param tf.Tensor strings: strings which are being scores. white-space delimited words.
    :return: same shape as `strings`, float32
    :rtype: tf.Tensor
    """
    return get_tf_mod().ken_lm_abs_score_strings(handle=handle, strings=strings)


def ken_lm_abs_score_bpe_strings(handle, bpe_merge_symbol, strings):
    """
    :param tf.Tensor handle: TF resource handle returned by :func:`ken_lm_load`
    :param str bpe_merge_symbol: e.g. "@@"
    :param tf.Tensor strings: strings which are being scores. white-space delimited words.
    :return: same shape as `strings`, float32
    :rtype: tf.Tensor
    """
    return get_tf_mod().ken_lm_abs_score_bpe_strings(handle=handle, bpe_merge_symbol=bpe_merge_symbol, strings=strings)


def ken_lm_abs_score_bpe_strings_dense(handle, bpe_merge_symbol, strings, labels):
    """
    :param tf.Tensor handle: TF resource handle returned by :func:`ken_lm_load`
    :param str bpe_merge_symbol: e.g. "@@"
    :param tf.Tensor strings: strings which are being scores. white-space delimited words.
    :param tf.Tensor|tf.Variable labels:
    :return: same shape as `strings`, float32
    :rtype: tf.Tensor
    """
    return get_tf_mod().ken_lm_abs_score_bpe_strings_dense(
        handle=handle, bpe_merge_symbol=bpe_merge_symbol, strings=strings, labels=labels
    )


if __name__ == "__main__":
    from returnn.util import better_exchook

    better_exchook.install()
    # Try to compile now.
    get_tf_mod(verbose=True)
    # Some demo.
    input_strings = sys.argv[1:] or ["hello world </s>"]
    test_lm_file = kenlm_dir + "/lm/test.arpa"
    assert os.path.exists(test_lm_file)
    lm_tf = ken_lm_load(filename=test_lm_file)
    input_strings_tf = tf.compat.v1.placeholder(tf.string, [None])
    output_scores_tf = ken_lm_abs_score_strings(handle=lm_tf, strings=input_strings_tf)
    with tf.compat.v1.Session() as session:
        output_scores = session.run(output_scores_tf, feed_dict={input_strings_tf: input_strings})
        print("input strings:", input_strings, "(sys.argv[1:])")
        print("output scores:", output_scores)
