// fstext/context-fst.cc

// Copyright      2018  Johns Hopkins University (author: Daniel Povey)

// See ../../COPYING for clarification regarding multiple authors
//
// 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
//
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
// MERCHANTABLITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.

#include "fstext/context-fst.h"
#include "base/kaldi-error.h"

namespace fst {


InverseContextFst::InverseContextFst(
    Label subsequential_symbol,
    const vector<int32>& phones,
    const vector<int32>& disambig_syms,
    int32 context_width,
    int32 central_position):
    context_width_(context_width),
    central_position_(central_position),
    phone_syms_(phones),
    disambig_syms_(disambig_syms),
    subsequential_symbol_(subsequential_symbol) {

  {  // This block checks the inputs.
    KALDI_ASSERT(subsequential_symbol != 0
                 && disambig_syms_.count(subsequential_symbol) == 0
                 && phone_syms_.count(subsequential_symbol) == 0);
    if (phone_syms_.empty())
      KALDI_WARN << "Context FST created but there are no phone symbols: probably "
          "input FST was empty.";
    KALDI_ASSERT(phone_syms_.count(0) == 0 && disambig_syms_.count(0) == 0 &&
                 central_position_ >= 0 && central_position_ < context_width_);
    for (size_t i = 0; i < phones.size(); i++) {
      KALDI_ASSERT(disambig_syms_.count(phones[i]) == 0);
    }
  }

  // empty vector, will be the ilabel_info vector that corresponds to epsilon,
  // in case our FST needs to output epsilons.
  vector<int32> empty_vec;
  Label epsilon_label = FindLabel(empty_vec);

  // epsilon_vec is the phonetic context window we have at the very start of a
  // sequence, meaning "no real phones have been seen yet".
  vector<int32> epsilon_vec(context_width_ - 1, 0);
  StateId start_state = FindState(epsilon_vec);

  KALDI_ASSERT(epsilon_label == 0 && start_state == 0);

  if (context_width_ > central_position_ + 1 && !disambig_syms_.empty()) {
    // We add a symbol whose sequence representation is [ 0 ], and whose
    // symbol-id is 1.  This is treated as a disambiguation symbol, we call it
    // #-1 in printed form.  It is necessary to ensure that all determinizable
    // LG's will have determinizable CLG's.  The problem it fixes is quite
    // subtle-- it relates to reordering of disambiguation symbols (they appear
    // earlier in CLG than in LG, relative to phones), and the fact that if a
    // disambig symbol appears at the very start of a sequence in CLG, it's not
    // clear exatly where it appeared on the corresponding sequence at the input
    // of LG.
    vector<int32> pseudo_eps_vec;
    pseudo_eps_vec.push_back(0);
    pseudo_eps_symbol_= FindLabel(pseudo_eps_vec);
    KALDI_ASSERT(pseudo_eps_symbol_ == 1);
  } else {
    pseudo_eps_symbol_ = 0;  // use actual epsilon.
  }
}


void InverseContextFst::ShiftSequenceLeft(Label label,
                                          std::vector<int32> *phone_seq) {
  if (!phone_seq->empty()) {
    phone_seq->erase(phone_seq->begin());
    phone_seq->push_back(label);
  }
}

void InverseContextFst::GetFullPhoneSequence(
    const std::vector<int32> &seq, Label label,
    std::vector<int32> *full_phone_sequence) {
  int32 context_width = context_width_;
  full_phone_sequence->reserve(context_width);
  full_phone_sequence->insert(full_phone_sequence->end(),
                              seq.begin(), seq.end());
  full_phone_sequence->push_back(label);
  for (int32 i = central_position_ + 1; i < context_width; i++) {
    if ((*full_phone_sequence)[i] == subsequential_symbol_) {
      (*full_phone_sequence)[i] = 0;
    }
  }
}


InverseContextFst::Weight InverseContextFst::Final(StateId s) {
  KALDI_ASSERT(static_cast<size_t>(s) < state_seqs_.size());

  const vector<int32> &phone_context = state_seqs_[s];

  KALDI_ASSERT(phone_context.size() == context_width_ - 1);

  bool has_final_prob;

  if (central_position_ < context_width_ - 1) {
    has_final_prob = (phone_context[central_position_] == subsequential_symbol_);
    // if phone_context[central_position_] != subsequential_symbol_ then we have
    // pending phones-in-context that we still need to output, so we need to
    // consume more subsequential symbols before we can terminate.
  } else {
    has_final_prob = true;
  }
  return has_final_prob ? Weight::One() : Weight::Zero();
}

bool InverseContextFst::GetArc(StateId s, Label ilabel, Arc *arc) {
  KALDI_ASSERT(ilabel != 0 && static_cast<size_t>(s) < state_seqs_.size() &&
               state_seqs_[s].size() == context_width_ - 1);

  if (IsDisambigSymbol(ilabel)) {
    // A disambiguation-symbol self-loop arc.
    CreateDisambigArc(s, ilabel, arc);
    return true;
  } else if (IsPhoneSymbol(ilabel)) {
    const vector<int32> &seq = state_seqs_[s];
    if (!seq.empty() && seq.back() == subsequential_symbol_) {
      return false;  // A real phone is not allowed to follow the subsequential
                     // symbol.
    }

    // next_seq will be 'seq' shifted left by 1, with 'ilabel' appended.
    vector<int32> next_seq(seq);
    ShiftSequenceLeft(ilabel, &next_seq);

    // full-seq will be the full context window of size context_width_.
    vector<int32> full_seq;
    GetFullPhoneSequence(seq, ilabel, &full_seq);

    StateId next_s = FindState(next_seq);

    CreatePhoneOrEpsArc(s, next_s, ilabel, full_seq, arc);
    return true;
  } else if (ilabel == subsequential_symbol_) {
    const vector<int32> &seq = state_seqs_[s];

    if (central_position_ + 1 == context_width_ ||
        seq[central_position_] == subsequential_symbol_) {
      // We already had "enough" subsequential symbols in a row and don't want to
      // accept any more, or we'd be making the subsequential symbol the central phone.
      return false;
    }

    // full-seq will be the full context window of size context_width_.
    vector<int32> full_seq;
    GetFullPhoneSequence(seq, ilabel, &full_seq);

    vector<int32> next_seq(seq);
    ShiftSequenceLeft(ilabel, &next_seq);
    StateId next_s = FindState(next_seq);

    CreatePhoneOrEpsArc(s, next_s, ilabel, full_seq, arc);
    return true;
  } else {
    KALDI_ERR << "ContextFst: CreateArc, invalid ilabel supplied [confusion "
              << "about phone list or disambig symbols?]: " << ilabel;
  }
  return false;  // won't get here.  suppress compiler error.
}


void InverseContextFst::CreateDisambigArc(StateId s, Label ilabel, Arc *arc) {
  // Creates a self-loop arc corresponding to the disambiguation symbol.
  vector<int32> label_info;       // This will be a vector containing just [ -olabel ].
  label_info.push_back(-ilabel);  // olabel is a disambiguation symbol.  Use its negative
                                  // so we can more easily distinguish them from phones.
  Label olabel = FindLabel(label_info);
  arc->ilabel = ilabel;
  arc->olabel = olabel;
  arc->weight = Weight::One();
  arc->nextstate = s;  // self-loop.
}

void InverseContextFst::CreatePhoneOrEpsArc(StateId src, StateId dest,
                                            Label ilabel,
                                            const vector<int32> &phone_seq,
                                            Arc *arc) {
  KALDI_PARANOID_ASSERT(phone_seq[central_position_] != subsequential_symbol_);

  arc->ilabel = ilabel;
  arc->weight = Weight::One();
  arc->nextstate = dest;
  if (phone_seq[central_position_] == 0) {
    // This can happen at the beginning of the graph.  In this case we don't
    // output a real phone, we createdt an epsilon arc (but sometimes we need to
    // use a special disambiguation symbol instead of epsilon).
    arc->olabel = pseudo_eps_symbol_;
  } else {
    // We have a phone in the central position.
    arc->olabel = FindLabel(phone_seq);
  }
}

StdArc::StateId InverseContextFst::FindState(const vector<int32> &seq) {
  // Finds state-id corresponding to this vector of phones.  Inserts it if
  // necessary.
  KALDI_ASSERT(static_cast<int32>(seq.size()) == context_width_ - 1);
  VectorToStateMap::const_iterator iter = state_map_.find(seq);
  if (iter == state_map_.end()) {  // Not already in map.
    StateId this_state_id = (StateId)state_seqs_.size();
    state_seqs_.push_back(seq);
    state_map_[seq] = this_state_id;
    return this_state_id;
  } else {
    return iter->second;
  }
}

StdArc::Label InverseContextFst::FindLabel(const vector<int32> &label_vec) {
  // Finds the ilabel corresponding to this vector (creates a new ilabel if
  // necessary).
  VectorToLabelMap::const_iterator iter = ilabel_map_.find(label_vec);
  if (iter == ilabel_map_.end()) {  // Not already in map.
    Label this_label = ilabel_info_.size();
    ilabel_info_.push_back(label_vec);
    ilabel_map_[label_vec] = this_label;
    return this_label;
  } else {
    return iter->second;
  }
}


void ComposeContext(const vector<int32> &disambig_syms_in,
                    int32 context_width, int32 central_position,
                    VectorFst<StdArc> *ifst,
                    VectorFst<StdArc> *ofst,
                    vector<vector<int32> > *ilabels_out,
                    bool project_ifst) {
  KALDI_ASSERT(ifst != NULL && ofst != NULL);
  KALDI_ASSERT(context_width > 0);
  KALDI_ASSERT(central_position >= 0);
  KALDI_ASSERT(central_position < context_width);

  vector<int32> disambig_syms(disambig_syms_in);
  std::sort(disambig_syms.begin(), disambig_syms.end());

  vector<int32> all_syms;
  GetInputSymbols(*ifst, false/*no eps*/, &all_syms);
  std::sort(all_syms.begin(), all_syms.end());
  vector<int32> phones;
  for (size_t i = 0; i < all_syms.size(); i++)
    if (!std::binary_search(disambig_syms.begin(),
                            disambig_syms.end(), all_syms[i]))
      phones.push_back(all_syms[i]);

  // Get subsequential symbol that does not clash with
  // any disambiguation symbol or symbol in the FST.
  int32 subseq_sym = 1;
  if (!all_syms.empty())
    subseq_sym = std::max(subseq_sym, all_syms.back() + 1);
  if (!disambig_syms.empty())
    subseq_sym = std::max(subseq_sym, disambig_syms.back() + 1);

  // if central_position == context_width-1, it's left-context, and no
  // subsequential symbol is needed.
  if (central_position != context_width-1) {
    AddSubsequentialLoop(subseq_sym, ifst);
    if (project_ifst) {
      fst::Project(ifst, fst::PROJECT_INPUT);
    }
  }

  InverseContextFst inv_c(subseq_sym, phones, disambig_syms,
                          context_width, central_position);

  // The following statement is equivalent to the following
  // (if FSTs had the '*' operator for composition):
  //   (*ofst) = inv(inv_c) * (*ifst)
  ComposeDeterministicOnDemandInverse(*ifst, &inv_c, ofst);

  inv_c.SwapIlabelInfo(ilabels_out);
}

void AddSubsequentialLoop(StdArc::Label subseq_symbol,
                          MutableFst<StdArc> *fst) {
  typedef StdArc Arc;
  typedef typename Arc::StateId StateId;
  typedef typename Arc::Weight Weight;

  vector<StateId> final_states;
  for (StateIterator<MutableFst<Arc> > siter(*fst); !siter.Done(); siter.Next()) {
    StateId s = siter.Value();
    if (fst->Final(s) != Weight::Zero())  final_states.push_back(s);
  }

  StateId superfinal = fst->AddState();
  Arc arc(subseq_symbol, 0, Weight::One(), superfinal);
  fst->AddArc(superfinal, arc);  // loop at superfinal.
  fst->SetFinal(superfinal, Weight::One());

  for (size_t i = 0; i < final_states.size(); i++) {
    StateId s = final_states[i];
    fst->AddArc(s, Arc(subseq_symbol, 0, fst->Final(s), superfinal));
    // No, don't remove the final-weights of the original states..
    // this is so we can add the subsequential loop in cases where
    // there is no context, and it won't hurt.
    // fst->SetFinal(s, Weight::Zero());
    arc.nextstate = final_states[i];
  }
}

void WriteILabelInfo(std::ostream &os, bool binary,
                     const vector<vector<int32> > &info) {
  int32 size = info.size();
  kaldi::WriteBasicType(os, binary, size);
  for (int32 i = 0; i < size; i++) {
    kaldi::WriteIntegerVector(os, binary, info[i]);
  }
}


void ReadILabelInfo(std::istream &is, bool binary,
                    vector<vector<int32> > *info) {
  int32 size = info->size();
  kaldi::ReadBasicType(is, binary, &size);
  info->resize(size);
  for (int32 i = 0; i < size; i++) {
    kaldi::ReadIntegerVector(is, binary, &((*info)[i]));
  }
}

SymbolTable *CreateILabelInfoSymbolTable(const vector<vector<int32> > &info,
                                         const SymbolTable &phones_symtab,
                                         std::string separator,
                                         std::string initial_disambig) {  // e.g. separator = "/", initial-disambig="#-1"
  KALDI_ASSERT(!info.empty() && info[0].empty());
  SymbolTable *ans = new SymbolTable("ilabel-info-symtab");
  int64 s = ans->AddSymbol(phones_symtab.Find(static_cast<int64>(0)));
  assert(s == 0);
  for (size_t i = 1; i < info.size(); i++) {
    if (info[i].size() == 0) {
      KALDI_ERR << "Invalid ilabel-info";
    }
    if (info[i].size() == 1 &&
       info[i][0] <= 0) {
      if (info[i][0] == 0) {  // special symbol at start that we want to call #-1.
        s = ans->AddSymbol(initial_disambig);
        if (s != i) {
          KALDI_ERR << "Disambig symbol " << initial_disambig
                    << " already in vocab";
        }
      } else {
        std::string disambig_sym = phones_symtab.Find(-info[i][0]);
        if (disambig_sym == "") {
          KALDI_ERR << "Disambig symbol " << -info[i][0]
                    << " not in phone symbol-table";
        }
        s = ans->AddSymbol(disambig_sym);
        if (s != i) {
          KALDI_ERR << "Disambig symbol " << disambig_sym
                    << " already in vocab";
        }
      }
    } else {
      // is a phone-context-window.
      std::string newsym;
      for (size_t j = 0; j < info[i].size(); j++) {
        std::string phonesym = phones_symtab.Find(info[i][j]);
        if (phonesym == "") {
          KALDI_ERR << "Symbol " << info[i][j]
                    << " not in phone symbol-table";
        }
        if (j != 0) newsym += separator;
        newsym += phonesym;
      }
      int64 s = ans->AddSymbol(newsym);
      if (s != static_cast<int64>(i)) {
        KALDI_ERR << "Some problem with duplicate symbols";
      }
    }
  }
  return ans;
}




}  // end namespace fst