#include "generate_model_forms.h"

#ifndef MAIN_TEST

#include "guessify_fst.h"

#include <algorithm>
#include <fstream>

using hfst::HfstOneLevelPaths;
using hfst::HfstOneLevelPath;

using hfst::TROPICAL_OPENFST_TYPE;

InvalidModelLine::InvalidModelLine(const std::string &line):
  line(line)
{}

std::ostream &operator<<(std::ostream &out,const StringVector &v)
{
  for (StringVector::const_iterator it = v.begin();
       it != v.end();
       ++it)
    {
      out << *it;
    }

  return out;
}

HfstTokenizer get_alphabet_string_tokenizer(HfstTransducer &fst)
{
  // FIXME: temporary until optimized lookup transducers implement
  // get_alphabet.
  HfstTransducer temp(fst);
  temp.convert(TROPICAL_OPENFST_TYPE);

  StringSet alphabet = temp.get_alphabet();
  
  HfstTokenizer tokenizer;

  for (StringSet::const_iterator it = alphabet.begin();
       it != alphabet.end();
       ++it)
    { tokenizer.add_multichar_symbol(*it); }

  return tokenizer;
}

StringVector get_analysis_prefix(const StringVector &reversed_analysis)
{
  StringVector prefix;

  // We want a prefix of a reversed string, so we iterate from end to
  // beginning.
  for (StringVector::const_reverse_iterator it = reversed_analysis.rbegin();
       it != reversed_analysis.rend();
       ++it)
    {
      prefix.push_back(*it);
      
      if (is_cathegory_symbol(*it))
    {
      // When we return, we have to reverse the result.
      return StringVector(prefix.rbegin(), prefix.rend());
    }
    }
  
  // It should actually be impossible to get here, since valid
  // analyses contain at least one cathegory symbol.
  return prefix;
}

StringVector join(StringVector sv1, const StringVector &sv2)
{
  sv1.insert(sv1.end(), sv2.begin(), sv2.end());
  return sv1;
}

bool contains_analysis_symbols(const StringVector &word_form)
{
  for (StringVector::const_iterator it = word_form.begin();
       it != word_form.end();
       ++it)
    {
      const std::string &symbol = *it;
      if (symbol.size() > 1 and
      symbol[0] == '[' and
      symbol[symbol.size() - 1] == ']')
    { return true; }
    }
  
  return false;
}

StringVector generate_word_forms(const StringVector &analysis,
                                 HfstTransducer &form_generator,
                                 size_t max_generated_forms,
                                 float generate_threshold)
{
  HfstOneLevelPaths * word_forms = form_generator.lookup(analysis);

  StringVectorSet result_set;

  size_t num = 1;
  
  float best_weight = -1;

  for (HfstOneLevelPaths::const_iterator it = word_forms->begin();
       it != word_forms->end();
       ++it)
    {
      if (num > max_generated_forms)
        { break; }

      if (best_weight == -1)
        { best_weight = it->first; }

      if (it->first - best_weight >= generate_threshold)
        { break; }

      const StringVector &word_form = it->second;
      
      if (contains_analysis_symbols(word_form))
        { continue; }
      
      // The word form is reversed, so we start from the end and
      // iterate to the beginning.
      result_set.insert(StringVector(word_form.rbegin(),
                                     word_form.rend()));
      
      ++num;
    }

  delete word_forms;

  StringVector results;

  bool first_form = true;

  for (StringVectorSet::const_iterator it = result_set.begin();
       it != result_set.end();
       ++it)
    {
      if (not first_form)
        { results.push_back(", "); }

      results.insert(results.end(), it->begin(), it->end());

      first_form = false;
    }

  if (results.empty())
    { results.push_back("<no word forms>"); }

  return results;
}

StringVectorVector get_model_forms(const StringVector &reversed_analysis,
                                   const StringVectorVector &model_forms,
                                   HfstTransducer &form_generator,
                                   size_t max_generated_forms,
                                   float generate_threshold)
{
  StringVector reversed_analysis_prefix =
    get_analysis_prefix(reversed_analysis);

  StringVectorVector results;

  for (StringVectorVector::const_iterator it  = model_forms.begin();
       it != model_forms.end();
       ++it)
    {
      StringVector model_analysis = join(*it, reversed_analysis_prefix);

      results.push_back(generate_word_forms(model_analysis,
                                            form_generator,
                                            max_generated_forms,
                                            generate_threshold));
    }

  return results;
}

StringPair split(const std::string &line,const std::string &separator)
{
  size_t separator_pos = line.find(separator);
  
  return StringPair(line.substr(0,separator_pos),
            line.substr(separator_pos + 1));
}

StringVector read_model_form(std::istream &in, HfstTokenizer &tokenizer)
{
  std::string line;
  std::getline(in,line);

  if (line.find(MODEL_FORM_PREFIX) != 0)
    { throw InvalidModelLine(line); }

  std::string model_form = line.substr(std::string(MODEL_FORM_PREFIX).size());

  StringVector tokenized_model_form = tokenizer.tokenize_one_level(model_form);
  std::reverse(tokenized_model_form.begin(),tokenized_model_form.end());

  return tokenized_model_form;
}

StringVectorVector read_model_forms(const std::string &model_form_filename,
                    HfstTokenizer &tokenizer)
{
  std::ifstream in(model_form_filename.c_str());

  if (in.peek() == EOF)
    { throw InvalidModelFile(); }

  StringVectorVector results;
  
  while (in.peek() != EOF)
    {
      results.push_back(read_model_form(in, tokenizer));
    }
  
  return results;
}


StringVectorVector get_guesses(const std::string &word_form,
                   HfstTransducer &guesser,
                   size_t number_of_guesses,
                   HfstTokenizer &tokenizer)
{
  StringVector tokenized_line = tokenizer.tokenize_one_level(word_form);
  std::reverse(tokenized_line.begin(),tokenized_line.end());

  HfstOneLevelPaths * paths = guesser.lookup_fd(tokenized_line);

  size_t num = 1;

  StringVectorVector results;

  for (HfstOneLevelPaths::const_iterator it = paths->begin();
       it != paths->end();
       ++it)
    {
      if (num > number_of_guesses)
    { break; }
      results.push_back(it->second);
      ++num;
    }

  delete paths;

  return results;
}

StringVectorVector get_paradigms(const std::string &word_form,
                                 const StringVectorVector &guesses,
                                 HfstTransducer &generator,
                                 const StringVectorVector &model_forms,
                                 size_t number_of_generated_forms,
                                 float generate_threshold)
{
  StringVectorVector paradigm_guesses;

  for (StringVectorVector::const_iterator it = guesses.begin();
       it != guesses.end();
       ++it)
    {
      StringVector analysis_guess = *it;

      StringVectorVector results = get_model_forms(analysis_guess,
                                                   model_forms,
                                                   generator,
                                                   number_of_generated_forms,
                                                   generate_threshold);
      
      StringVector paradigm;
      paradigm.push_back(word_form);
      paradigm.push_back("\t");

      StringVector rev_analysis_guess(analysis_guess);
      std::reverse(rev_analysis_guess.begin(), rev_analysis_guess.end());
      paradigm.insert(paradigm.end(),
              rev_analysis_guess.begin(),
              rev_analysis_guess.end());

      for (StringVectorVector::const_iterator jt = results.begin();
       jt != results.end();
       ++jt)
        {
          const StringVector &model_form = *jt;

          paradigm.push_back("\t");
          paradigm.insert(paradigm.end(),
                  model_form.begin(),
                  model_form.end());
        }
      
      paradigm_guesses.push_back(paradigm);
    }

  return paradigm_guesses;
}

#else // MAIN_TEST

#include <cstdio>
#include <iostream>
#include <exception>

using hfst::HfstInputStream;

int main(int argc, char * argv[])
{
  // Open input fst-file for reading the guesser and generator.
  std::cerr << "Reading guesser and generator from file "
        << argv[1] << std::endl;

  HfstInputStream in(argv[1]);

  HfstTransducer guesser(in);
  HfstTransducer generator(in);

  // Use the generator alphabet (same as the guesser alphabet) to
  // create a tokenizer.
  HfstTokenizer tokenizer = get_alphabet_string_tokenizer(generator);

  // Read the model analyses, that are used to generate word forms.
  std::cerr << "Reading model analyses from file " << argv[2] << std::endl;

  StringVectorVector model_forms = read_model_forms(argv[2],tokenizer);

  // Read input words and generate model forms. Input words are
  // tokenized and reversed for guessing. The MAX_ANALYSES first
  // analysis guesses for each word are split into a prefix consisting
  // of the lemma, pos-tag and categroy symbol. -each of the prefixes
  // is combined with each of the model analyses and the generator is
  // used to generate the corresponding word-forms, which are
  // displayed for the user.
  std::cerr << "Reading unknown word forms from STDIN." << std::endl;
  while (std::cin.peek() != EOF)
    {
      std::string line;
      std::getline(std::cin,line);
      
      StringVectorVector guesses = get_guesses(line,
                           guesser,
                           MAX_ANALYSES,
                           tokenizer);

      StringVectorVector guess_model_forms = get_paradigms(line,
                               guesses,
                               generator,
                               model_forms,
                               MAX_GENERATED_FORMS)
    ;
      

      for (StringVectorVector::const_iterator it = guess_model_forms.begin();
       it != guess_model_forms.end();
       ++it)
    {
      std::cout << *it << std::endl;
    }
    }
}
#endif // MAIN_TEST

