// nnet3/discriminative-supervision.cc
  
  // Copyright 2012-2015  Johns Hopkins University (author: Daniel Povey)
  //           2014-2015  Vimal Manohar
  
  // 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 "nnet3/discriminative-supervision.h"
  #include "lat/lattice-functions.h"
  
  namespace kaldi {
  namespace discriminative {
  
  
  DiscriminativeSupervision::DiscriminativeSupervision(
      const DiscriminativeSupervision &other):
      weight(other.weight), num_sequences(other.num_sequences),
      frames_per_sequence(other.frames_per_sequence),
      num_ali(other.num_ali), den_lat(other.den_lat) { }
  
  void DiscriminativeSupervision::Swap(DiscriminativeSupervision *other) {
    std::swap(weight, other->weight);
    std::swap(num_sequences, other->num_sequences);
    std::swap(frames_per_sequence, other->frames_per_sequence);
    std::swap(num_ali, other->num_ali);
    std::swap(den_lat, other->den_lat);
  }
  
  bool DiscriminativeSupervision::operator == (
      const DiscriminativeSupervision &other) const {
    return ( weight == other.weight &&
        num_sequences == other.num_sequences &&
        frames_per_sequence == other.frames_per_sequence &&
        num_ali == other.num_ali &&
        fst::Equal(den_lat, other.den_lat) );
  }
  
  void DiscriminativeSupervision::Write(std::ostream &os, bool binary) const {
    WriteToken(os, binary, "<DiscriminativeSupervision>");
    WriteToken(os, binary, "<Weight>");
    WriteBasicType(os, binary, weight);
    WriteToken(os, binary, "<NumSequences>");
    WriteBasicType(os, binary, num_sequences);
    WriteToken(os, binary, "<FramesPerSeq>");
    WriteBasicType(os, binary, frames_per_sequence);
    KALDI_ASSERT(frames_per_sequence > 0 &&
                 num_sequences > 0);
  
    WriteToken(os, binary, "<NumAli>");
    WriteIntegerVector(os, binary, num_ali);
  
    WriteToken(os, binary, "<DenLat>");
    if (!WriteLattice(os, binary, den_lat)) {
      // We can't return error status from this function so we
      // throw an exception.
      KALDI_ERR << "Error writing denominator lattice to stream";
    }
  
    WriteToken(os, binary, "</DiscriminativeSupervision>");
  }
  
  void DiscriminativeSupervision::Read(std::istream &is, bool binary) {
    ExpectToken(is, binary, "<DiscriminativeSupervision>");
    ExpectToken(is, binary, "<Weight>");
    ReadBasicType(is, binary, &weight);
    ExpectToken(is, binary, "<NumSequences>");
    ReadBasicType(is, binary, &num_sequences);
    ExpectToken(is, binary, "<FramesPerSeq>");
    ReadBasicType(is, binary, &frames_per_sequence);
    KALDI_ASSERT(frames_per_sequence > 0 &&
                 num_sequences > 0);
  
    ExpectToken(is, binary, "<NumAli>");
    ReadIntegerVector(is, binary, &num_ali);
  
    ExpectToken(is, binary, "<DenLat>");
    {
      Lattice *lat = NULL;
      if (!ReadLattice(is, binary, &lat) || lat == NULL) {
        // We can't return error status from this function so we
        // throw an exception.
        KALDI_ERR << "Error reading Lattice from stream";
      }
      den_lat = *lat;
      delete lat;
      TopSort(&den_lat);
    }
  
    ExpectToken(is, binary, "</DiscriminativeSupervision>");
  }
  
  bool DiscriminativeSupervision::Initialize(const std::vector<int32> &num_ali,
                                             const Lattice &den_lat,
                                             BaseFloat weight) {
    if (num_ali.size() == 0) return false;
    if (den_lat.NumStates() == 0) return false;
  
    this->weight = weight;
    this->num_sequences = 1;
    this->frames_per_sequence = num_ali.size();
    this->num_ali = num_ali;
    this->den_lat = den_lat;
    KALDI_ASSERT(TopSort(&(this->den_lat)));
  
    // Checks if num frames in alignment matches lattice
    Check();
  
    return true;
  }
  
  void DiscriminativeSupervision::Check() const {
    int32 num_frames_subsampled = num_ali.size();
    KALDI_ASSERT(num_frames_subsampled ==
                 num_sequences * frames_per_sequence);
  
    {
      std::vector<int32> state_times;
      int32 max_time = LatticeStateTimes(den_lat, &state_times);
      KALDI_ASSERT(max_time == num_frames_subsampled);
    }
  }
  
  DiscriminativeSupervisionSplitter::DiscriminativeSupervisionSplitter(
      const SplitDiscriminativeSupervisionOptions &config,
      const TransitionModel &tmodel,
      const DiscriminativeSupervision &supervision):
      config_(config), tmodel_(tmodel), supervision_(supervision) {
    if (supervision_.num_sequences != 1) {
      KALDI_WARN << "Splitting already-reattached sequence (only expected in "
                 << "testing code)";
    }
  
    KALDI_ASSERT(supervision_.num_sequences == 1); // For now, don't allow splitting already merged examples
  
    den_lat_ = supervision_.den_lat;
    PrepareLattice(&den_lat_, &den_lat_scores_);
  
    int32 num_states = den_lat_.NumStates(),
          num_frames = supervision_.frames_per_sequence * supervision_.num_sequences;
    KALDI_ASSERT(num_states > 0);
    int32 start_state = den_lat_.Start();
    // Lattice should be top-sorted and connected, so start-state must be 0.
    KALDI_ASSERT(start_state == 0 && "Expecting start-state to be 0");
  
    KALDI_ASSERT(num_states == den_lat_scores_.state_times.size());
    KALDI_ASSERT(den_lat_scores_.state_times[start_state] == 0);
    KALDI_ASSERT(den_lat_scores_.state_times.back() == num_frames);
  }
  
  // Make sure that for any given pdf-id and any given frame, the den-lat has
  // only one transition-id mapping to that pdf-id, on the same frame.
  // It helps us to more completely minimize the lattice.  Note: we
  // can't do this if the criterion is MPFE, because in that case the
  // objective function will be affected by the phone-identities being
  // different even if the pdf-ids are the same.
  void DiscriminativeSupervisionSplitter::CollapseTransitionIds(
      const std::vector<int32> &state_times, Lattice *lat) const {
    typedef Lattice::StateId StateId;
    typedef Lattice::Arc Arc;
  
    int32 num_frames = state_times.back();   // TODO: Check if this is always true
    StateId num_states = lat->NumStates();
  
    std::vector<std::map<int32, int32> > pdf_to_tid(num_frames);
    for (StateId s = 0; s < num_states; s++) {
      int32 t = state_times[s];
      for (fst::MutableArcIterator<Lattice> aiter(lat, s);
           !aiter.Done(); aiter.Next()) {
        KALDI_ASSERT(t >= 0 && t < num_frames);
        Arc arc = aiter.Value();
        KALDI_ASSERT(arc.ilabel != 0 && arc.ilabel == arc.olabel);
        int32 pdf = tmodel_.TransitionIdToPdf(arc.ilabel);
        if (pdf_to_tid[t].count(pdf) != 0) {
          arc.ilabel = arc.olabel = pdf_to_tid[t][pdf];
          aiter.SetValue(arc);
        } else {
          pdf_to_tid[t][pdf] = arc.ilabel;
        }
      }
    }
  }
  
  void DiscriminativeSupervisionSplitter::LatticeInfo::Check() const {
    // Check if all the vectors are of size num_states
    KALDI_ASSERT(state_times.size() == alpha.size() &&
                 state_times.size() == beta.size());
  
    // Check that the states are ordered in increasing order of state_times.
    // This must be true since the states are in breadth-first search order.
    KALDI_ASSERT(IsSorted(state_times));
  }
  
  void DiscriminativeSupervisionSplitter::GetFrameRange(int32 begin_frame, int32 num_frames, bool normalize,
                                                        DiscriminativeSupervision *out_supervision) const {
    int32 end_frame = begin_frame + num_frames;
    // Note: end_frame is not included in the range of frames that the
    // output supervision object covers; it's one past the end.
    KALDI_ASSERT(num_frames > 0 && begin_frame >= 0 &&
                 begin_frame + num_frames <=
                 supervision_.num_sequences * supervision_.frames_per_sequence);
  
    CreateRangeLattice(den_lat_,
                       den_lat_scores_,
                       begin_frame, end_frame, normalize,
                       &(out_supervision->den_lat));
  
    out_supervision->num_ali.clear();
    std::copy(supervision_.num_ali.begin() + begin_frame,
              supervision_.num_ali.begin() + end_frame,
              std::back_inserter(out_supervision->num_ali));
  
    out_supervision->num_sequences = 1;
    out_supervision->weight = supervision_.weight;
    out_supervision->frames_per_sequence = num_frames;
  
    out_supervision->Check();
  }
  
  void DiscriminativeSupervisionSplitter::CreateRangeLattice(
      const Lattice &in_lat, const LatticeInfo &scores,
      int32 begin_frame, int32 end_frame, bool normalize,
      Lattice *out_lat) const {
    typedef Lattice::StateId StateId;
  
    const std::vector<int32> &state_times = scores.state_times;
  
    // Some checks to ensure the lattice and scores are prepared properly
    KALDI_ASSERT(state_times.size() == in_lat.NumStates());
    if (!in_lat.Properties(fst::kTopSorted, true))
      KALDI_ERR << "Input lattice must be topologically sorted.";
  
    std::vector<int32>::const_iterator begin_iter =
        std::lower_bound(state_times.begin(), state_times.end(), begin_frame),
        end_iter = std::lower_bound(begin_iter,
                                    state_times.end(), end_frame);
  
    KALDI_ASSERT(*begin_iter == begin_frame &&
                 (begin_iter == state_times.begin() ||
                  begin_iter[-1] < begin_frame));
    // even if end_frame == supervision_.num_frames, there should be a state with
    // that frame index.
    KALDI_ASSERT(end_iter[-1] < end_frame &&
                 (end_iter < state_times.end() || *end_iter == end_frame));
    StateId begin_state = begin_iter - state_times.begin(),
            end_state = end_iter - state_times.begin();
  
    KALDI_ASSERT(end_state > begin_state);
    out_lat->DeleteStates();
    out_lat->ReserveStates(end_state - begin_state + 2);
  
    // Add special start state
    StateId start_state = out_lat->AddState();
    out_lat->SetStart(start_state);
  
    for (StateId i = begin_state; i < end_state; i++)
      out_lat->AddState();
  
    // Add the special final-state.
    StateId final_state = out_lat->AddState();
    out_lat->SetFinal(final_state, LatticeWeight::One());
  
    for (StateId state = begin_state; state < end_state; state++) {
      StateId output_state = state - begin_state + 1;
      if (state_times[state] == begin_frame) {
        // we'd like to make this an initial state, but OpenFst doesn't allow
        // multiple initial states.  Instead we add an epsilon transition to it
        // from our actual initial state.  The weight on this
        // transition is the forward probability of the said 'initial state'
        LatticeWeight weight = LatticeWeight::One();
        weight.SetValue1((normalize ? scores.beta[0] : 0.0) - scores.alpha[state]);
        // Add negative of the forward log-probability to the graph cost score,
        // since the acoustic scores would be changed later.
        // Assuming that the lattice is scaled with appropriate acoustic
        // scale.
        // We additionally normalize using the total lattice score. Since the
        // same score is added as normalizer to all the paths in the lattice,
        // the relative probabilities of the paths in the lattice is not affected.
        // Note: Doing a forward-backward on this split must result in a total
        // score of 0 because of the normalization.
  
        out_lat->AddArc(start_state,
                        LatticeArc(0, 0, weight, output_state));
      } else {
        KALDI_ASSERT(scores.state_times[state] < end_frame);
      }
      for (fst::ArcIterator<Lattice> aiter(in_lat, state);
            !aiter.Done(); aiter.Next()) {
        const LatticeArc &arc = aiter.Value();
        StateId nextstate = arc.nextstate;
        if (nextstate >= end_state) {
          // A transition to any state outside the range becomes a transition to
          // our special final-state.
          // The weight is just the negative of the backward log-probability +
          // the arc cost. We again normalize with the total lattice score.
          LatticeWeight weight;
          //KALDI_ASSERT(scores.beta[state] < 0);
          weight.SetValue1(arc.weight.Value1() - scores.beta[nextstate]);
          weight.SetValue2(arc.weight.Value2());
          // Add negative of the backward log-probability to the LM score, since
          // the acoustic scores would be changed later.
          // Note: We don't normalize here because that is already done with the
          // initial cost.
  
          out_lat->AddArc(output_state,
              LatticeArc(arc.ilabel, arc.olabel, weight, final_state));
        } else {
          StateId output_nextstate = nextstate - begin_state + 1;
          out_lat->AddArc(output_state,
              LatticeArc(arc.ilabel, arc.olabel, arc.weight, output_nextstate));
        }
      }
    }
  
    // Get rid of the word labels and put the
    // transition-ids on both sides.
    fst::Project(out_lat, fst::PROJECT_INPUT);
    fst::RmEpsilon(out_lat);
  
    if (config_.collapse_transition_ids)
      CollapseTransitionIds(state_times, out_lat);
  
    if (config_.determinize) {
      if (!config_.minimize) {
        Lattice tmp_lat;
        fst::Determinize(*out_lat, &tmp_lat);
        std::swap(*out_lat, tmp_lat);
      } else {
        Lattice tmp_lat;
        fst::Reverse(*out_lat, &tmp_lat);
        fst::Determinize(tmp_lat, out_lat);
        fst::Reverse(*out_lat, &tmp_lat);
        fst::Determinize(tmp_lat, out_lat);
        fst::RmEpsilon(out_lat);
      }
    }
  
    fst::TopSort(out_lat);
    std::vector<int32> state_times_tmp;
    KALDI_ASSERT(LatticeStateTimes(*out_lat, &state_times_tmp) ==
                                              end_frame - begin_frame);
  
    // Remove the acoustic scale that was previously added
    if (config_.acoustic_scale != 1.0) {
      fst::ScaleLattice(fst::AcousticLatticeScale(
            1 / config_.acoustic_scale), out_lat);
    }
  }
  
  void DiscriminativeSupervisionSplitter::PrepareLattice(
      Lattice *lat, LatticeInfo *scores) const {
    // Scale the lattice to appropriate acoustic scale. It is important to
    // ensure this is equal to the acoustic scale used while training. This is
    // because, on splitting lattices, the initial and final costs are added
    // into the graph cost.
    KALDI_ASSERT(config_.acoustic_scale != 0.0);
    if (config_.acoustic_scale != 1.0)
      fst::ScaleLattice(fst::AcousticLatticeScale(
          config_.acoustic_scale), lat);
  
    LatticeStateTimes(*lat, &(scores->state_times));
    int32 num_states = lat->NumStates();
    std::vector<std::pair<int32,int32> > state_time_indexes(num_states);
    for (int32 s = 0; s < num_states; s++) {
      state_time_indexes[s] = std::make_pair(scores->state_times[s], s);
    }
  
    // Order the states based on the state times. This is stronger than just
    // topological sort. This is required by the lattice splitting code.
    std::sort(state_time_indexes.begin(), state_time_indexes.end());
  
    std::vector<int32> state_order(num_states);
    for (int32 s = 0; s < num_states; s++) {
      state_order[state_time_indexes[s].second] = s;
    }
  
    fst::StateSort(lat, state_order);
    ComputeLatticeScores(*lat, scores);
  }
  
  void DiscriminativeSupervisionSplitter::ComputeLatticeScores(const Lattice &lat,
      LatticeInfo *scores) const {
    LatticeStateTimes(lat, &(scores->state_times));
    ComputeLatticeAlphasAndBetas(lat, false,
                                 &(scores->alpha), &(scores->beta));
    scores->Check();
    // This check will fail if the lattice is not breadth-first search sorted
  }
  
  void MergeSupervision(const std::vector<const DiscriminativeSupervision*> &input,
      DiscriminativeSupervision *output_supervision) {
    KALDI_ASSERT(!input.empty());
    int32 num_inputs = input.size();
    if (num_inputs == 1) {
      *output_supervision = *(input[0]);
      return;
    }
    *output_supervision = *(input[num_inputs-1]);
    for (int32 i = num_inputs - 2; i >= 0; i--) {
      const DiscriminativeSupervision &src = *(input[i]);
      KALDI_ASSERT(src.num_sequences == 1);
      if (output_supervision->weight == src.weight &&
          output_supervision->frames_per_sequence ==
          src.frames_per_sequence) {
        // Combine with current output
        // append src.den_lat to output_supervision->den_lat.
        fst::Concat(src.den_lat, &output_supervision->den_lat);
  
        output_supervision->num_ali.insert(
            output_supervision->num_ali.begin(),
            src.num_ali.begin(), src.num_ali.end());
  
        output_supervision->num_sequences++;
      } else {
        KALDI_ERR << "Mismatch weight or frames_per_sequence  between inputs";
      }
    }
    DiscriminativeSupervision &out_sup = *output_supervision;
    fst::TopSort(&(out_sup.den_lat));
    out_sup.Check();
  }
  
  } // namespace discriminative
  } // namespace kaldi