/*==========================================================================
 * Copyright (c) 2002-2011 University of Massachusetts.  All Rights Reserved.
 *
 * Use of the Lemur Toolkit for Language Modeling and Information Retrieval
 * is subject to the terms of the software license set forth in the LICENSE
 * file included with this software, and also available at
 * http://www.lemurproject.org/license.html
 *
 *==========================================================================
*/
/*
 * < Calculation of Field-level Collection Statistics >
 *                        author : jykim@cs.umass.edu
 * 
 */

#include "lemur/common_headers.hpp" 
#include "indri/Repository.hpp"
#include "indri/CompressedCollection.hpp"
#include "indri/LocalQueryServer.hpp"
#include "indri/ScopedLock.hpp"
#include "indri/QueryEnvironment.hpp"
#include "indri/Index.hpp"
#include <iostream>
using namespace std;

// Fetch the list of words from the collection
void getVocabulary(indri::collection::Repository& r,map<string,int>& wordList,int& documentCount);

// 
void getFieldLM(indri::collection::Repository& r,string fieldName,ofstream& output,map<string,int>& wordList,int documentCount);;

int main(int argc, char *argv[]) {
  indri::collection::Repository r;
  string repName=argv[1];
  string fileName = argv[2];
  //string 
  r.openRead(repName);
  cout<<"Repository open"<<endl;

  map<string,int> wordList;
  int documentCount;
  getVocabulary(r,wordList,documentCount);
  /*
  ofstream output_wordlist("wordList.txt");
  for(map<string,int>::iterator iter=wordList.begin();iter!=wordList.end();iter++){
    output_wordlist<<iter->first<<" "<<iter->second<<endl;
  }
  output_wordlist.close();
  cout<<documentCount<<endl;
  */
  ofstream output(fileName.c_str());

  indri::server::LocalQueryServer local(r);
  vector<string> fields = local.fieldList();
  for( size_t i=0; i<fields.size(); i++ ) {
    std::cout << fields[i] << " ";
    getFieldLM(r,fields[i],output,wordList,documentCount);
  }
  std::cout << std::endl;
  output.close();

  return 0;

}

// Fetch the list of words from the collection
void getVocabulary(indri::collection::Repository& r,map<string,int>& wordList,int& documentCount){
  indri::collection::Repository::index_state state = r.indexes();

  indri::index::Index* index = (*state)[0];
  indri::index::VocabularyIterator* iter = index->vocabularyIterator();

  iter->startIteration();
  std::cout << "UniqueTermCount: " << index->uniqueTermCount() << "DocumentCount: " << index->documentCount() << std::endl;

  documentCount=index->documentCount();

  int id=1;

  while( !iter->finished() ) {
    indri::index::DiskTermData* entry = iter->currentEntry();
    indri::index::TermData* termData = entry->termData;

    if(wordList.count(termData->term)!=0){
      cerr<<"term has existed!!"<<endl;
      exit(-1);
    }
    wordList[termData->term]=id;
    id++;

    iter->nextEntry();
  }
  cout<<"id:"<<id<<endl;

  delete iter;
}

// 
void getFieldLM(indri::collection::Repository& r,string fieldName,ofstream& output,map<string,int>& wordList,int documentCount){


  map<string,int> fieldLM;
  for(map<string,int>::iterator iter=wordList.begin();iter!=wordList.end();iter++){
    fieldLM[iter->first] = 0;
  }
  double sum=0;
  
  indri::server::LocalQueryServer local(r);
  	
  std::vector<lemur::api::DOCID_T> documentIDs;
  for(int i=1;i<=documentCount;i++)
    documentIDs.push_back(i);

  indri::server::QueryServerVectorsResponse* response = local.documentVectors( documentIDs );
  
  for(int i=0;i<response->getResults().size();i++){
    indri::api::DocumentVector* docVector = response->getResults()[i];
          
    for( size_t j=0; j<docVector->fields().size(); j++ ) {
      const indri::api::DocumentVector::Field& field = docVector->fields()[j];
      if(field.name==fieldName){
        int begin=field.begin;
        int end=field.end;
        for(int k=begin;k<end;k++){
          int stem_pos=docVector->positions()[k];
          string stem=docVector->stems()[stem_pos];
          if(stem=="[OOV]")
            continue;
          if(wordList.count(stem)==0){
            cerr<<"stem:"<<stem<<" not found."<<endl;
            exit(-1);
          }
          fieldLM[stem]++;
          sum++;
        }
      }
    }
    delete docVector;
  }
	
  //output language model
  output<<fieldName<<"\t";
  for(map<string,int>::iterator iter=fieldLM.begin();iter!=fieldLM.end();iter++){
    if(iter->second > 0)
      output<<iter->first<<" "<<iter->second<<"\t";///sum
  }
  output<<endl;
	
}