// Copyright 2016 Google Inc.

//

// 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.!

#ifndef MODEL_INTERFACE_H_

#define MODEL_INTERFACE_H_

#include <memory>

#include <set>

#include <string>

#include <utility>

#include <vector>

#include "common.h"

#include "normalizer.h"

#include "sentencepiece_model.pb.h"

#include "sentencepiece_processor.h"

#include "third_party/absl/container/flat_hash_map.h"

#include "third_party/absl/strings/string_view.h"

#include "third_party/darts_clone/darts.h"

#include "util.h"

namespace sentencepiece {

// "_this_is_a_pen" => ["_this", "_is", "_a", "_pen"]

std::vector<absl::string_view> SplitIntoWords(

    absl::string_view text, bool treat_ws_as_suffix = false,

    bool allow_ws_only_pieces = false);

// Converts byte (0-255) to piece (e.g., 58 -> "<0x3A>").

std::string ByteToPiece(unsigned char c);

// Converts piece to byte (e.g., "<0x3A>" -> 58). Returns -1 if `piece` is not

// a valid byte piece.

int PieceToByte(absl::string_view piece);

using EncodeResult = std::vector<std::pair<absl::string_view, int>>;

using NBestEncodeResult = std::vector<std::pair<EncodeResult, float>>;

class ModelProto;

// Underlying model interface.

// Given a normalized string, returns a sequence of sentence pieces with ids.

class ModelInterface {

 public:

  using PieceToIdMap = absl::flat_hash_map<absl::string_view, int>;

  //                                           string_util::string_view_hash>;

  absl::string_view unk_piece() const;

  absl::string_view bos_piece() const;

  absl::string_view eos_piece() const;

  absl::string_view pad_piece() const;

  // `model_proto` should not be deleted until ModelInterface is destroyed.

  explicit ModelInterface(const ModelProto &model_proto);

  ModelInterface() {}

  virtual ~ModelInterface();

  // Returns Status.

  // Encode/Decode functions are valid only when status is OK.

  virtual util::Status status() const { return status_; }

  virtual const ModelProto &model_proto() const { return *model_proto_; }

  virtual const normalizer::PrefixMatcher *prefix_matcher() const {

    return matcher_.get();

  }

  // Given a normalized string, returns a sequence of sentence pieces with ids.

  // The concatenation of pieces must be the same as `normalized`.

  virtual EncodeResult Encode(absl::string_view normalized) const = 0;

  // The same as above, but returns nbest result with score.

  virtual NBestEncodeResult NBestEncode(absl::string_view normalized,

                                        int nbest_size) const {

    LOG(ERROR) << "Not implemented.";

    return NBestEncodeResult();

  }

  virtual EncodeResult SampleEncode(absl::string_view normalized,

                                    float alpha) const {

    LOG(ERROR) << "Not implemented.";

    return EncodeResult();

  }

  // Sample `samples` many tokenisations from the segmentation lattice

  // If `wor` is true, the samples are taken without replacement, and the scores

  // are the inclusion probabilities of the elements in the sample; otherwise

  // the samples are taken with replacement and the scores are the log-probs of

  // sample elements

  // If `include_best` is true, the best tokenisation is always included in the

  // sample, and the remaining elements are sampled excluding the best.

  virtual NBestEncodeResult SampleEncodeAndScore(absl::string_view normalized,

                                                 float alpha, int samples,

                                                 bool wor,

                                                 bool include_best) const {

    LOG(ERROR) << "Not implemented.";

    return {{EncodeResult(), 0.0}};

  }

  // Calculates the entropy of the segmentation lattice with inverse temperature

  // `alpha`. Uses a novel dynamic program to calculate the entropy.

  virtual float CalculateEntropy(absl::string_view normalized,

                                 float alpha) const {

    LOG(ERROR) << "Not implemented.";

    return 0.0;

  }

  // Return true if SampleEncode returns a valid result.

  virtual bool IsSampleEncodeAvailable() const { return false; }

  // Return true if NBestEncode returns a valid result.

  virtual bool IsNBestEncodeAvailable() const { return false; }

  // Return true if SampleEncodeAndScore returns a valid result.

  virtual bool IsSampleEncodeAndScoreAvailable() const { return false; }

  // Return true if CalculateEntropy returns a valid result.

  virtual bool IsCalculateEntropyAvailable() const { return false; }

  // Returns the vocab id of `piece`.

  // Returns UNK(0) if `piece` is unknown

  virtual int PieceToId(absl::string_view piece) const;

  // Returns the string representation of vocab with `id`.

  // id must be 0 <= id < GetPieceSize().

  virtual const std::string &IdToPiece(int id) const {

    return model_proto_->pieces(id).piece();

  }

  // Returns the size of sentence pieces, which is the same

  // as the size of vocabulary for NMT.

  virtual int GetPieceSize() const {

    if (!model_proto_) return 0;

    return model_proto_->pieces_size();

  }

  // Returns the score of `id`.

  // Score represents a log probability of the piece.

  // We can roughly estimate the unigram frequency of the piece.

  virtual float GetScore(int id) const {

    return model_proto_->pieces(id).score();

  }

  // Returns true if `id` is unknown symbol.

  virtual bool IsUnknown(int id) const {

    return (model_proto_->pieces(id).type() ==

            ModelProto::SentencePiece::UNKNOWN);

  }

  // Returns true if `id` is control symbol.

  virtual bool IsControl(int id) const {

    return (model_proto_->pieces(id).type() ==

            ModelProto::SentencePiece::CONTROL);

  }

  // Returns true if `id` is unused symbol.

  virtual bool IsUnused(int id) const {

    return (model_proto_->pieces(id).type() ==

            ModelProto::SentencePiece::UNUSED);

  }

  // Returns true if `id` is user defined symbol.

  virtual bool IsUserDefined(int id) const {

    return (model_proto_->pieces(id).type() ==

            ModelProto::SentencePiece::USER_DEFINED);

  }

  // Returns true if `id` is byte symbol.

  virtual bool IsByte(int id) const {

    return (model_proto_->pieces(id).type() == ModelProto::SentencePiece::BYTE);

  }

  virtual bool ByteFallbackEnabled() const {

    return model_proto_ && model_proto_->trainer_spec().byte_fallback();

  }

  // Verifies if the `expected` and `actual` outputs are equivalent. `expected`

  // and `actual` are sentence pieces joined by space (` `). Normally it means

  // that the two strings are identical. In some model, due to float rounding

  // errors, the strings may not be identical, but they may be still equivalent

  // provided their scores are close enough (by some espilon).

  virtual bool VerifyOutputsEquivalent(absl::string_view expected,

                                       absl::string_view actual) const {

    return expected == actual;

  }

 protected:

  void InitializePieces();

  // Non-virtual (inlined) implementation for faster execution.

  inline float GetScoreInlined(int id) const {

    return model_proto_->pieces(id).score();

  }

  inline bool IsUnknownInlined(int id) const {

    return (model_proto_->pieces(id).type() ==

            ModelProto::SentencePiece::UNKNOWN);

  }

  inline bool IsControlInlined(int id) const {

    return (model_proto_->pieces(id).type() ==

            ModelProto::SentencePiece::CONTROL);

  }

  inline bool IsUnusedInlined(int id) const {

    return (model_proto_->pieces(id).type() ==

            ModelProto::SentencePiece::UNUSED);

  }

  inline bool IsUserDefinedInlined(int id) const {

    return (model_proto_->pieces(id).type() ==

            ModelProto::SentencePiece::USER_DEFINED);

  }

  inline bool IsByteInlined(int id) const {

    return (model_proto_->pieces(id).type() == ModelProto::SentencePiece::BYTE);

  }

  const ModelProto *model_proto_ = nullptr;

  // PrefixMatcher for user defined symbols.

  std::unique_ptr<normalizer::PrefixMatcher> matcher_;

  // piece -> id map for normal pieces

  PieceToIdMap pieces_;

  // piece -> id map for control, unknown, and byte pieces

  PieceToIdMap reserved_id_map_;

  // unknown id.

  int unk_id_ = 0;

  // status.

  util::Status status_;

};

}  // namespace sentencepiece

#endif  // MODEL_INTERFACE_H_