// chain/chain-den-graph.cc
  
  // Copyright      2015-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 "chain/chain-den-graph.h"
  #include "hmm/hmm-utils.h"
  #include "fstext/push-special.h"
  
  namespace kaldi {
  namespace chain {
  
  
  DenominatorGraph::DenominatorGraph(const fst::StdVectorFst &fst,
                                     int32 num_pdfs):
      num_pdfs_(num_pdfs) {
    SetTransitions(fst, num_pdfs);
    SetInitialProbs(fst);
  }
  
  const Int32Pair* DenominatorGraph::BackwardTransitions() const {
    return backward_transitions_.Data();
  }
  
  const Int32Pair* DenominatorGraph::ForwardTransitions() const {
    return forward_transitions_.Data();
  }
  
  const DenominatorGraphTransition* DenominatorGraph::Transitions() const {
    return transitions_.Data();
  }
  
  const CuVector<BaseFloat>& DenominatorGraph::InitialProbs() const {
    return initial_probs_;
  }
  
  void DenominatorGraph::SetTransitions(const fst::StdVectorFst &fst,
                                        int32 num_pdfs) {
    int32 num_states = fst.NumStates();
  
    std::vector<std::vector<DenominatorGraphTransition> >
        transitions_out(num_states),
        transitions_in(num_states);
    for (int32 s = 0; s < num_states; s++) {
      for (fst::ArcIterator<fst::StdVectorFst> aiter(fst, s); !aiter.Done();
           aiter.Next()) {
        const fst::StdArc &arc = aiter.Value();
        DenominatorGraphTransition transition;
        transition.transition_prob = exp(-arc.weight.Value());
        transition.pdf_id = arc.ilabel - 1;
        transition.hmm_state = arc.nextstate;
        KALDI_ASSERT(transition.pdf_id >= 0 && transition.pdf_id < num_pdfs);
        transitions_out[s].push_back(transition);
        // now the reverse transition.
        transition.hmm_state = s;
        transitions_in[arc.nextstate].push_back(transition);
      }
    }
  
    std::vector<Int32Pair> forward_transitions(num_states);
    std::vector<Int32Pair> backward_transitions(num_states);
    std::vector<DenominatorGraphTransition> transitions;
  
    for (int32 s = 0; s < num_states; s++) {
      forward_transitions[s].first = static_cast<int32>(transitions.size());
      transitions.insert(transitions.end(), transitions_out[s].begin(),
                         transitions_out[s].end());
      forward_transitions[s].second = static_cast<int32>(transitions.size());
    }
    for (int32 s = 0; s < num_states; s++) {
      backward_transitions[s].first = static_cast<int32>(transitions.size());
      transitions.insert(transitions.end(), transitions_in[s].begin(),
                         transitions_in[s].end());
      backward_transitions[s].second = static_cast<int32>(transitions.size());
    }
  
    forward_transitions_ = forward_transitions;
    backward_transitions_ = backward_transitions;
    transitions_ = transitions;
  }
  
  void DenominatorGraph::SetInitialProbs(const fst::StdVectorFst &fst) {
    // we set only the start-state to have probability mass, and then 100
    // iterations of HMM propagation, over which we average the probabilities.
    // initial probs won't end up making a huge difference as we won't be using
    // derivatives from the first few frames, so this isn't 100% critical.
    int32 num_iters = 100;
    int32 num_states = fst.NumStates();
  
    // we normalize each state so that it sums to one (including
    // final-probs)... this is needed because the 'chain' code doesn't
    // have transition probabilities.
    Vector<double> normalizing_factor(num_states);
    for (int32 s = 0; s < num_states; s++) {
      double tot_prob = exp(-fst.Final(s).Value());
      for (fst::ArcIterator<fst::StdVectorFst> aiter(fst, s); !aiter.Done();
           aiter.Next()) {
        tot_prob += exp(-aiter.Value().weight.Value());
      }
      KALDI_ASSERT(tot_prob > 0.0 && tot_prob < 100.0);
      normalizing_factor(s) = 1.0 / tot_prob;
    }
  
    Vector<double> cur_prob(num_states), next_prob(num_states),
        avg_prob(num_states);
    cur_prob(fst.Start()) = 1.0;
    for (int32 iter = 0; iter < num_iters; iter++) {
      avg_prob.AddVec(1.0 / num_iters, cur_prob);
      for (int32 s = 0; s < num_states; s++) {
        double prob = cur_prob(s) * normalizing_factor(s);
  
        for (fst::ArcIterator<fst::StdVectorFst> aiter(fst, s); !aiter.Done();
             aiter.Next()) {
          const fst::StdArc &arc = aiter.Value();
          next_prob(arc.nextstate) += prob * exp(-arc.weight.Value());
        }
      }
      cur_prob.Swap(&next_prob);
      next_prob.SetZero();
      // Renormalize, beause the HMM won't sum to one even after the
      // previous normalization (due to final-probs).
      cur_prob.Scale(1.0 / cur_prob.Sum());
    }
  
    Vector<BaseFloat> avg_prob_float(avg_prob);
    initial_probs_ = avg_prob_float;
  }
  
  void DenominatorGraph::GetNormalizationFst(const fst::StdVectorFst &ifst,
                                             fst::StdVectorFst *ofst) {
    KALDI_ASSERT(ifst.NumStates() == initial_probs_.Dim());
    if (&ifst != ofst)
      *ofst = ifst;
    int32 new_initial_state = ofst->AddState();
    Vector<BaseFloat> initial_probs(initial_probs_);
  
    for (int32 s = 0; s < initial_probs_.Dim(); s++) {
      BaseFloat initial_prob = initial_probs(s);
      KALDI_ASSERT(initial_prob > 0.0);
      fst::StdArc arc(0, 0, fst::TropicalWeight(-log(initial_prob)), s);
      ofst->AddArc(new_initial_state, arc);
      ofst->SetFinal(s, fst::TropicalWeight::One());
    }
    ofst->SetStart(new_initial_state);
    fst::RmEpsilon(ofst);
    fst::ArcSort(ofst, fst::ILabelCompare<fst::StdArc>());
  }
  
  
  void MapFstToPdfIdsPlusOne(const TransitionModel &trans_model,
                             fst::StdVectorFst *fst) {
    int32 num_states = fst->NumStates();
    for (int32 s = 0; s < num_states; s++) {
      for (fst::MutableArcIterator<fst::StdVectorFst> aiter(fst, s);
           !aiter.Done(); aiter.Next()) {
        fst::StdArc arc = aiter.Value();
        KALDI_ASSERT(arc.ilabel == arc.olabel);
        if (arc.ilabel > 0) {
          arc.ilabel = trans_model.TransitionIdToPdf(arc.ilabel) + 1;
          arc.olabel = arc.ilabel;
          aiter.SetValue(arc);
        }
      }
    }
  }
  
  void MinimizeAcceptorNoPush(fst::StdVectorFst *fst) {
    BaseFloat delta = fst::kDelta * 10.0;  // use fairly loose delta for
                                           // aggressive minimimization.
    fst::ArcMap(fst, fst::QuantizeMapper<fst::StdArc>(delta));
    fst::EncodeMapper<fst::StdArc> encoder(fst::kEncodeLabels | fst::kEncodeWeights,
                                           fst::ENCODE);
    fst::Encode(fst, &encoder);
    fst::internal::AcceptorMinimize(fst);
    fst::Decode(fst, encoder);
  }
  
  // This static function, used in CreateDenominatorFst, sorts an
  // fst's states in decreasing order of number of transitions (into + out of)
  // the state.  The aim is to have states that have a lot of transitions
  // either into them or out of them, be numbered earlier, so hopefully
  // they will be scheduled first and won't delay the computation
  static void SortOnTransitionCount(fst::StdVectorFst *fst) {
    // negative_num_transitions[i] will contain (before sorting), the pair
    // ( -(num-transitions-into(i) + num-transition-out-of(i)), i)
    int32 num_states = fst->NumStates();
    std::vector<std::pair<int32, int32> > negative_num_transitions(num_states);
    for (int32 i = 0; i < num_states; i++) {
      negative_num_transitions[i].first = 0;
      negative_num_transitions[i].second = i;
    }
    for (int32 i = 0; i < num_states; i++) {
      for (fst::ArcIterator<fst::StdVectorFst> aiter(*fst, i); !aiter.Done();
           aiter.Next()) {
        negative_num_transitions[i].first--;
        negative_num_transitions[aiter.Value().nextstate].first--;
      }
    }
    std::sort(negative_num_transitions.begin(), negative_num_transitions.end());
    std::vector<fst::StdArc::StateId> order(num_states);
    for (int32 i = 0; i < num_states; i++)
      order[negative_num_transitions[i].second] = i;
    fst::StateSort(fst, order);
  }
  
  void DenGraphMinimizeWrapper(fst::StdVectorFst *fst) {
    for (int32 i = 1; i <= 3; i++) {
      fst::StdVectorFst fst_reversed;
      fst::Reverse(*fst, &fst_reversed);
      fst::PushSpecial(&fst_reversed, fst::kDelta * 0.01);
      MinimizeAcceptorNoPush(&fst_reversed);
      fst::Reverse(fst_reversed, fst);
      KALDI_LOG << "Number of states and arcs in transition-id FST after reversed "
                << "minimization is " << fst->NumStates() << " and "
                << NumArcs(*fst) << " (pass " << i << ")";
      fst::PushSpecial(fst, fst::kDelta * 0.01);
      MinimizeAcceptorNoPush(fst);
      KALDI_LOG << "Number of states and arcs in transition-id FST after regular "
                << "minimization is " << fst->NumStates() << " and "
                << NumArcs(*fst) << " (pass " << i << ")";
    }
    fst::RmEpsilon(fst);
    KALDI_LOG << "Number of states and arcs in transition-id FST after "
              << "removing any epsilons introduced by reversal is "
              << fst->NumStates() << " and "
              << NumArcs(*fst);
    fst::PushSpecial(fst, fst::kDelta * 0.01);
  }
  
  
  static void PrintDenGraphStats(const fst::StdVectorFst &den_graph) {
    int32 num_states = den_graph.NumStates();
    int32 degree_cutoff = 3;  // track states with <= transitions in/out.
    int32 num_states_low_degree_in = 0,
        num_states_low_degree_out = 0,
        tot_arcs = 0;
    std::vector<int32> num_in_arcs(num_states, 0);
    for (int32 s = 0; s < num_states; s++) {
      if (den_graph.NumArcs(s) <= degree_cutoff) {
        num_states_low_degree_out++;
      }
      tot_arcs += den_graph.NumArcs(s);
      for (fst::ArcIterator<fst::StdVectorFst> aiter(den_graph, s);
           !aiter.Done(); aiter.Next()) {
        int32 dest_state = aiter.Value().nextstate;
        num_in_arcs[dest_state]++;
      }
    }
    for (int32 s = 0; s < num_states; s++) {
      if (num_in_arcs[s] <= degree_cutoff) {
        num_states_low_degree_in++;
      }
    }
    KALDI_LOG << "Number of states is " << num_states << " and arcs "
              << tot_arcs << "; number of states with in-degree <= "
              << degree_cutoff << " is " << num_states_low_degree_in
              << " and with out-degree <= " << degree_cutoff
              << " is " << num_states_low_degree_out;
  }
  
  
  // Check that every pdf is seen, warn if some are not.
  static void CheckDenominatorFst(int32 num_pdfs,
                                  const fst::StdVectorFst &den_fst) {
    std::vector<bool> pdf_seen(num_pdfs);
    int32 num_states = den_fst.NumStates();
    for (int32 s = 0; s < num_states; s++) {
      for (fst::ArcIterator<fst::StdVectorFst> aiter(den_fst, s);
           !aiter.Done(); aiter.Next()) {
        int32 pdf_id = aiter.Value().ilabel - 1;
        KALDI_ASSERT(pdf_id >= 0 && pdf_id < num_pdfs);
        pdf_seen[pdf_id] = true;
      }
    }
    for (int32 pdf = 0; pdf < num_pdfs; pdf++) {
      if (!pdf_seen[pdf]) {
        KALDI_WARN << "Pdf-id " << pdf << " is not seen in denominator graph.";
      }
    }
  }
  
  void CreateDenominatorFst(const ContextDependency &ctx_dep,
                            const TransitionModel &trans_model,
                            const fst::StdVectorFst &phone_lm_in,
                            fst::StdVectorFst *den_fst) {
    using fst::StdVectorFst;
    using fst::StdArc;
    KALDI_ASSERT(phone_lm_in.NumStates() != 0);
    fst::StdVectorFst phone_lm(phone_lm_in);
  
    KALDI_LOG << "Number of states and arcs in phone-LM FST is "
              << phone_lm.NumStates() << " and " << NumArcs(phone_lm);
  
    int32 subsequential_symbol = trans_model.GetPhones().back() + 1;
    if (ctx_dep.CentralPosition() != ctx_dep.ContextWidth() - 1) {
      // note: this function only adds the subseq symbol to the input of what was
      // previously an acceptor, so we project, i.e. copy the ilabels to the
      // olabels
      AddSubsequentialLoop(subsequential_symbol, &phone_lm);
      fst::Project(&phone_lm, fst::PROJECT_INPUT);
    }
    std::vector<int32> disambig_syms;  // empty list of diambiguation symbols.
  
    // inv_cfst will be expanded on the fly, as needed.
    fst::InverseContextFst inv_cfst(subsequential_symbol,
                                    trans_model.GetPhones(),
                                    disambig_syms,
                                    ctx_dep.ContextWidth(),
                                    ctx_dep.CentralPosition());
  
    fst::StdVectorFst context_dep_lm;
    fst::ComposeDeterministicOnDemandInverse(phone_lm, &inv_cfst,
                                             &context_dep_lm);
  
    // at this point, context_dep_lm will have indexes into 'ilabels' as its
    // input symbol (representing context-dependent phones), and phones on its
    // output.  We don't need the phones, so we'll project.
    fst::Project(&context_dep_lm, fst::PROJECT_INPUT);
  
    KALDI_LOG << "Number of states and arcs in context-dependent LM FST is "
              << context_dep_lm.NumStates() << " and " << NumArcs(context_dep_lm);
  
    std::vector<int32> disambig_syms_h; // disambiguation symbols on input side
    // of H -- will be empty.
    HTransducerConfig h_config;
    // the default is 1, but just document that we want this to stay as one.
    // we'll use the same value in test time.  Consistency is the key here.
    h_config.transition_scale = 1.0;
  
    StdVectorFst *h_fst = GetHTransducer(inv_cfst.IlabelInfo(),
                                         ctx_dep,
                                         trans_model,
                                         h_config,
                                         &disambig_syms_h);
    KALDI_ASSERT(disambig_syms_h.empty());
    StdVectorFst transition_id_fst;
    TableCompose(*h_fst, context_dep_lm, &transition_id_fst);
    delete h_fst;
  
    BaseFloat self_loop_scale = 1.0;  // We have to be careful to use the same
                                      // value in test time.
    // 'reorder' must always be set to true for chain models.
    bool reorder = true;
    bool check_no_self_loops = true;
  
    // add self-loops to the FST with transition-ids as its labels.
    AddSelfLoops(trans_model, disambig_syms_h, self_loop_scale, reorder,
                 check_no_self_loops, &transition_id_fst);
    // at this point transition_id_fst will have transition-ids as its ilabels and
    // context-dependent phones (indexes into IlabelInfo()) as its olabels.
    // Discard the context-dependent phones by projecting on the input, keeping
    // only the transition-ids.
    fst::Project(&transition_id_fst, fst::PROJECT_INPUT);
  
    MapFstToPdfIdsPlusOne(trans_model, &transition_id_fst);
    KALDI_LOG << "Number of states and arcs in transition-id FST is "
              << transition_id_fst.NumStates() << " and "
              << NumArcs(transition_id_fst);
  
    // RemoveEpsLocal doesn't remove all epsilons, but it keeps the graph small.
    fst::RemoveEpsLocal(&transition_id_fst);
    // If there are remaining epsilons, remove them.
    fst::RmEpsilon(&transition_id_fst);
    KALDI_LOG << "Number of states and arcs in transition-id FST after "
              << "removing epsilons is "
              << transition_id_fst.NumStates() << " and "
              << NumArcs(transition_id_fst);
  
    DenGraphMinimizeWrapper(&transition_id_fst);
  
    SortOnTransitionCount(&transition_id_fst);
  
    *den_fst = transition_id_fst;
    CheckDenominatorFst(trans_model.NumPdfs(), *den_fst);
    PrintDenGraphStats(*den_fst);
  }
  
  
  int32 DenominatorGraph::NumStates() const {
    return forward_transitions_.Dim();
  }
  }  // namespace chain
  }  // namespace kaldi