#include "LRModel.h"
#include "moses/Range.h"
#include "moses/Bitmap.h"
#include "moses/InputType.h"
#include "HReorderingForwardState.h"
#include "HReorderingBackwardState.h"
#include "PhraseBasedReorderingState.h"
#include "BidirectionalReorderingState.h"
#include "SparseReordering.h"

namespace Moses
{

bool
IsMonotonicStep(Range  const& prev, // words range of last source phrase
                Range  const& cur,  // words range of current source phrase
                Bitmap const& cov)  // coverage bitmap
{
  size_t e = prev.GetEndPos() + 1;
  size_t s = cur.GetStartPos();
  return (s == e || (s >= e && !cov.GetValue(e)));
}

bool
IsSwap(Range const& prev, Range const& cur, Bitmap const& cov)
{
  size_t s = prev.GetStartPos();
  size_t e = cur.GetEndPos();
  return (e+1 == s || (e < s && !cov.GetValue(s-1)));
}

size_t
LRModel::
GetNumberOfTypes() const
{
  return ((m_modelType == MSD)  ? 3 :
          (m_modelType == MSLR) ? 4 : 2);
}

size_t
LRModel::
GetNumScoreComponents() const
{
  size_t score_per_dir = m_collapseScores ? 1 : GetNumberOfTypes();
  return ((m_direction == Bidirectional)
          ? 2 * score_per_dir + m_additionalScoreComponents
          : score_per_dir + m_additionalScoreComponents);
}

void
LRModel::
ConfigureSparse(const std::map<std::string,std::string>& sparseArgs,
                const LexicalReordering* producer)
{
  if (sparseArgs.size()) {
    m_sparse.reset(new SparseReordering(sparseArgs, producer));
  }
}

void
LRModel::
SetAdditionalScoreComponents(size_t number)
{
  m_additionalScoreComponents = number;
}

/// return orientation for the first phrase
LRModel::ReorderingType
LRModel::
GetOrientation(Range const& cur) const
{
  UTIL_THROW_IF2(m_modelType == None, "Reordering Model Type is None");
  return ((m_modelType == LeftRight) ? R :
          (cur.GetStartPos() == 0) ? M  :
          (m_modelType == MSD)     ? D  :
          (m_modelType == MSLR)    ? DR : NM);
}

LRModel::ReorderingType
LRModel::
GetOrientation(Range const& prev, Range const& cur) const
{
  UTIL_THROW_IF2(m_modelType == None, "No reordering model type specified");
  return ((m_modelType == LeftRight)
          ? prev.GetEndPos() <= cur.GetStartPos() ? R : L
        : (cur.GetStartPos() == prev.GetEndPos() + 1) ? M
          : (m_modelType == Monotonic) ? NM
          : (prev.GetStartPos() ==  cur.GetEndPos() + 1) ? S
          : (m_modelType == MSD) ? D
          : (cur.GetStartPos() > prev.GetEndPos()) ? DR : DL);
}

LRModel::ReorderingType
LRModel::
GetOrientation(int const reoDistance) const
{
  // this one is for HierarchicalReorderingBackwardState
  return ((m_modelType == LeftRight)
          ? (reoDistance >= 1) ? R : L
        : (reoDistance == 1) ? M
          : (m_modelType == Monotonic) ? NM
          : (reoDistance == -1)  ? S
          : (m_modelType == MSD) ? D
          : (reoDistance  >  1) ? DR : DL);
}

LRModel::ReorderingType
LRModel::
GetOrientation(Range const& prev, Range const& cur,
               Bitmap const& cov) const
{
  return ((m_modelType == LeftRight)
          ? cur.GetStartPos() > prev.GetEndPos() ? R : L
        : IsMonotonicStep(prev,cur,cov) ? M
          : (m_modelType == Monotonic) ? NM
          : IsSwap(prev,cur,cov) ? S
          : (m_modelType == MSD) ? D
          : cur.GetStartPos() > prev.GetEndPos() ? DR : DL);
}

LRModel::
LRModel(const std::string &modelType)
  : m_modelString(modelType)
  , m_scoreProducer(NULL)
  , m_modelType(None)
  , m_phraseBased(true)
  , m_collapseScores(false)
  , m_direction(Backward)
  , m_additionalScoreComponents(0)
{
  std::vector<std::string> config = Tokenize<std::string>(modelType, "-");

  for (size_t i=0; i<config.size(); ++i) {
    if      (config[i] == "hier")   {
      m_phraseBased = false;
    } else if (config[i] == "phrase") {
      m_phraseBased = true;
    } else if (config[i] == "wbe")    {
      m_phraseBased = true;
    }
    // no word-based decoding available, fall-back to phrase-based
    // This is the old lexical reordering model combination of moses

    else if (config[i] == "msd")          {
      m_modelType = MSD;
    } else if (config[i] == "mslr")         {
      m_modelType = MSLR;
    } else if (config[i] == "monotonicity") {
      m_modelType = Monotonic;
    } else if (config[i] == "leftright")    {
      m_modelType = LeftRight;
    }

    // unidirectional is deprecated, use backward instead
    else if (config[i] == "unidirectional") {
      m_direction = Backward;
    } else if (config[i] == "backward")       {
      m_direction = Backward;
    } else if (config[i] == "forward")        {
      m_direction = Forward;
    } else if (config[i] == "bidirectional")  {
      m_direction = Bidirectional;
    }

    else if (config[i] == "f")  {
      m_condition = F;
    } else if (config[i] == "fe") {
      m_condition = FE;
    }

    else if (config[i] == "collapseff") {
      m_collapseScores = true;
    } else if (config[i] == "allff") {
      m_collapseScores = false;
    } else {
      std::cerr
          << "Illegal part in the lexical reordering configuration string: "
          << config[i] << std::endl;
      exit(1);
    }
  }

  if (m_modelType == None) {
    std::cerr
        << "You need to specify the type of the reordering model "
        << "(msd, monotonicity,...)" << std::endl;
    exit(1);
  }
}

LRState *
LRModel::
CreateLRState(const InputType &input) const
{
  LRState *bwd = NULL, *fwd = NULL;
  size_t offset = 0;

  switch(m_direction) {
  case Backward:
  case Bidirectional:
    if (m_phraseBased)
      bwd = new PhraseBasedReorderingState(*this, Backward, offset);
    else
      bwd = new HReorderingBackwardState(*this, offset);
    offset += m_collapseScores ? 1 : GetNumberOfTypes();
    if (m_direction == Backward) return bwd; // else fall through
  case Forward:
    if (m_phraseBased)
      fwd = new PhraseBasedReorderingState(*this, Forward, offset);
    else
      fwd = new HReorderingForwardState(*this, input.GetSize(), offset);
    offset += m_collapseScores ? 1 : GetNumberOfTypes();
    if (m_direction == Forward) return fwd;
  }
  return new BidirectionalReorderingState(*this, bwd, fwd, 0);
}

}