package edu.stanford.nlp.patterns;

import edu.stanford.nlp.io.IOUtils;
import edu.stanford.nlp.ling.CoreLabel;
import edu.stanford.nlp.patterns.surface.Token;
import edu.stanford.nlp.pipeline.CoreNLPProtos;
import edu.stanford.nlp.pipeline.ProtobufAnnotationSerializer;
import edu.stanford.nlp.util.CollectionValuedMap;
import edu.stanford.nlp.util.LuceneFieldType;
import edu.stanford.nlp.util.ArgumentParser;
import edu.stanford.nlp.util.ArgumentParser.Option;
import edu.stanford.nlp.util.logging.Redwood;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.core.KeywordAnalyzer;
import org.apache.lucene.document.*;
import org.apache.lucene.index.*;
import org.apache.lucene.queryparser.classic.ParseException;
import org.apache.lucene.search.*;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.Version;

import java.io.*;

import java.util.*;
import java.util.function.Function;



/**
 * To create a lucene inverted index from tokens to sentence ids.
 * (Right now it is not storing all core tokens although some functions might suggest that.)
 *
 * @author Sonal Gupta on 10/14/14.
 */
public class LuceneSentenceIndex<E extends Pattern> extends SentenceIndex<E> {

  @Option(name="saveTokens")
  boolean saveTokens = false;

  IndexWriter indexWriter;
  File indexDir = null;
  Directory dir;

  Analyzer analyzer = new KeywordAnalyzer();
//  Analyzer analyzer = new Analyzer() {
//    @Override
//    protected TokenStreamComponents createComponents(String fieldName, Reader reader) {
//      Tokenizer source = new KeywordTokenizer(reader);
//      TokenStream result = new StopWordsFilter(source);
//      return new TokenStreamComponents(source, result);
//    }
//  };

//  public final class StopWordsFilter extends FilteringTokenFilter {
//    /**
//     * Build a filter that removes words that are too long or too
//     * short from the text.
//     */
//    public StopWordsFilter(TokenStream in) {
//      super(true, in);
//    }
//
//    @Override
//    public boolean accept() throws IOException {
//      return !stopWords.contains(input.toString().toLowerCase());
//    }
//  }

  //StandardAnalyzer analyzer = new StandardAnalyzer(Version.LUCENE_42);
  IndexWriterConfig iwc=new IndexWriterConfig(analyzer);
  DirectoryReader reader = null;
  IndexSearcher searcher;
  ProtobufAnnotationSerializer p = new ProtobufAnnotationSerializer();

  //The fields in index are: tokens, sentence id, List<CoreLabel> annotation of the sentence (optional; not used when sentences are in memory)
  public LuceneSentenceIndex(Properties props, Set<String> stopWords, String indexDirStr, Function<CoreLabel, Map<String, String>> transformer) {
    super(stopWords, transformer);
    ArgumentParser.fillOptions(this, props);
    indexDir = new File(indexDirStr);
  }


  void setIndexReaderSearcher() throws IOException {
    FSDirectory index = FSDirectory.open(indexDir.toPath());
    if(reader == null){
      reader = DirectoryReader.open(index);
      searcher = new IndexSearcher(reader);
    }else{
      DirectoryReader newreader = DirectoryReader.openIfChanged(reader);
      if(newreader != null) {
        reader.close();
        reader = newreader;
        searcher = new IndexSearcher(reader);
      }
    }
  }

//  SentenceIndex.SentenceIteratorWithWords queryIndex(SurfacePattern pat){
//
//
//    String[] n = pat.getSimplerTokensNext();
//    String[] pr = pat.getSimplerTokensPrev();
//    boolean rest = false;
//    if(n!=null){
//      for(String e: n){
//        if(!specialWords.contains(e)){
//          rest = true;
//          break;
//        }
//      }
//    }
//    if(rest == false && pr!=null){
//      for(String e: pr){
//        if(!specialWords.contains(e) && !stopWords.contains(e)){
//          rest = true;
//          break;
//        }
//      }
//    }
//
//  }

  /**
   * give all sentences that have these words
   * **/
  Set<String> queryIndexGetSentences(CollectionValuedMap<String, String> words) throws IOException, ParseException {
    setIndexReaderSearcher();
    BooleanQuery.Builder builder = new BooleanQuery.Builder();
    String pkey  = Token.getKeyForClass(PatternsAnnotations.ProcessedTextAnnotation.class);

    for(Map.Entry<String, Collection<String>> en: words.entrySet()){
      boolean processedKey = en.getKey().equals(pkey);
      for(String en2: en.getValue()){
          if(!processedKey || !stopWords.contains(en2.toLowerCase()))
            builder.add(new BooleanClause(new TermQuery(new Term(en.getKey(), en2)), BooleanClause.Occur.MUST));
      }
    }
    BooleanQuery query = builder.build();
    //query.add(new BooleanClause(new TermQuery(new Term("textannotation","sonal")), BooleanClause.Occur.MUST));

//    String queryStr = "";
//    for(Map.Entry<String, Collection<String>> en: words.entrySet()){
//      for(String en2: en.getValue()){
//        queryStr+= " " + en.getKey() + ":"+en2;
//      }
//    }
//    QueryParser queryParser = new QueryParser(Version.LUCENE_42, "sentence", analyzer);
//
//    queryParser.setDefaultOperator(QueryParser.Operator.AND);
//
//    Query query = queryParser.parse(queryStr);

    //Map<String, List<CoreLabel>> sents = null;
    TopDocs tp = searcher.search(query, Integer.MAX_VALUE);
    Set<String> sentids = new HashSet<>();
    if (tp.totalHits > 0) {
      for (ScoreDoc s : tp.scoreDocs) {
        int docId = s.doc;
        Document d = searcher.doc(docId);
//        byte[] sent = d.getBinaryValue("tokens").bytes;
//        if(saveTokens) {
//          sents = new HashMap<String, List<CoreLabel>>();
//          List<CoreLabel> tokens = readProtoBufAnnotation(sent);
//          sents.put(d.get("sentid"), tokens);
//        } else{
        sentids.add(d.get("sentid"));
        //}
      }
    } else {
      throw new RuntimeException("how come no documents for " + words + ". Query formed is " + query);
      //System.out.println("number of sentences for tokens " + words + " are " + sentids);
//    if(!saveTokens){
//      sents = getSentences(sentids);
//    }
    }
    return sentids;
  }


  @Override
  public void add(Map<String, DataInstance> sentences, boolean addProcessedText) {
    try {
      this.setIndexWriter();

      for(Map.Entry<String, DataInstance> en: sentences.entrySet()){
        //String sentence = StringUtils.joinWords(en.getValue(), " ");
        add(en.getValue().getTokens(), en.getKey(), addProcessedText);
      }

      indexWriter.commit();
      closeIndexWriter();

    } catch (IOException e) {
      throw new RuntimeException(e);
    }

  }

  @Override
  public Map<E, Set<String>> queryIndex(Collection<E> patterns) {
    try{
      Map<E, Set<String>> sents = new HashMap<>();
      for(E p : patterns){
        Set<String> sentids = queryIndexGetSentences(p.getRelevantWords());
        sents.put(p, sentids);
      }
      return sents;
    }
    catch (ParseException | IOException e) {
      throw new RuntimeException(e);
    }
  }



  public void listAllDocuments() throws IOException {
    setIndexReaderSearcher();
    for(int i = 0; i < reader.numDocs(); i++){
      Document d = searcher.doc(i);
//      byte[] sent = d.getBinaryValue("tokens").bytes;
//      List<CoreLabel> tokens = readProtoBufAnnotation(sent);
      System.out.println(d.get("sentid"));
    }
  }

  private List<CoreLabel> readProtoBufAnnotation(byte[] sent) throws IOException {
    ProtobufAnnotationSerializer p = new ProtobufAnnotationSerializer();
    List<CoreLabel> toks = new ArrayList<>();
    ByteArrayInputStream is = new ByteArrayInputStream(sent);
    CoreNLPProtos.Token d;
    do{
      d = CoreNLPProtos.Token.parseDelimitedFrom(is);
      if(d != null)
        toks.add(p.fromProto(d));
    } while(d != null);
    return toks;
  }



  byte[] getProtoBufAnnotation(List<CoreLabel> tokens) throws IOException {
    ByteArrayOutputStream os = new ByteArrayOutputStream();
    for(CoreLabel token: tokens){
      CoreNLPProtos.Token ptoken = p.toProto(token);
      ptoken.writeDelimitedTo(os);
    }
    os.flush();
    return os.toByteArray();
  }

  @Override
  protected void add(List<CoreLabel> tokens, String sentid, boolean addProcessedText){
    try{
      setIndexWriter();

      Document doc = new Document();

      for(CoreLabel l : tokens) {
        for (Map.Entry<String, String> en: transformCoreLabeltoString.apply(l).entrySet()) {
          doc.add(new StringField(en.getKey(), en.getValue(), Field.Store.YES));//, ANALYZED));
        }
        if(addProcessedText){
          String ptxt = l.get(PatternsAnnotations.ProcessedTextAnnotation.class);
          if(!stopWords.contains(ptxt.toLowerCase()))
            doc.add(new StringField(Token.getKeyForClass(PatternsAnnotations.ProcessedTextAnnotation.class), ptxt, Field.Store.YES));//, ANALYZED));
        }
      }
      doc.add(new StringField("sentid", sentid, Field.Store.YES));


      if(tokens!= null && saveTokens)
        doc.add(new Field("tokens", getProtoBufAnnotation(tokens), LuceneFieldType.NOT_INDEXED));

      indexWriter.addDocument(doc);

    }catch(IOException e){
      throw new RuntimeException(e);
    }
  }

  @Override
  public void finishUpdating() {
    if(indexWriter != null){
    try {
      indexWriter.commit();
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    }
    closeIndexWriter();
  }

  @Override
  public void update(List<CoreLabel> tokens, String sentid) {
    try {
      setIndexWriter();
      indexWriter.deleteDocuments(new TermQuery(new Term("sentid",sentid)));
      add(tokens, sentid, true);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }

  }

  void setIndexWriter()  {try{
    if(indexWriter == null){
      dir = FSDirectory.open(indexDir.toPath());
      Redwood.log(Redwood.DBG, "Updating lucene index at " + indexDir);
      indexWriter = new IndexWriter(dir, iwc);
    }}catch(IOException e){
    throw new RuntimeException(e);
  }
  }

  void closeIndexWriter(){
    try {
    if(indexWriter != null)
        indexWriter.close();
    indexWriter = null;

      if(dir != null)
      dir.close();
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }

  @Override
  public void saveIndex(String dir){
    if(!indexDir.toString().equals(dir)){
      try {
        IOUtils.cp(indexDir, new File(dir), true);
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
    }
  }


  public static LuceneSentenceIndex createIndex(Map<String, List<CoreLabel>> sentences,  Properties props, Set<String> stopWords, String indexDiskDir, Function<CoreLabel, Map<String, String>> transformer) {
    try{
      LuceneSentenceIndex sentindex = new LuceneSentenceIndex(props, stopWords, indexDiskDir, transformer);
      System.out.println("Creating lucene index at " + indexDiskDir);
      IOUtils.deleteDirRecursively(sentindex.indexDir);
      if(sentences!= null){
        sentindex.setIndexWriter();

        sentindex.add(sentences, true);

        sentindex.closeIndexWriter();
        sentindex.setIndexReaderSearcher();
        System.out.println("Number of documents added are " + sentindex.reader.numDocs());
        sentindex.numAllSentences += sentindex.reader.numDocs();
      }
      return sentindex;
    }catch(IOException e){
      throw new RuntimeException(e);
    }
  }

  public static LuceneSentenceIndex loadIndex(Properties props, Set<String> stopwords, String dir,  Function<CoreLabel, Map<String, String>> transformSentenceToString) {
    try {
    LuceneSentenceIndex sentindex = new LuceneSentenceIndex(props, stopwords, dir, transformSentenceToString);
    sentindex.setIndexReaderSearcher();
    System.out.println("Number of documents read from the index " + dir + " are " + sentindex.reader.numDocs());
    sentindex.numAllSentences += sentindex.reader.numDocs();
    return sentindex;
    } catch (IOException e) {
     throw new RuntimeException(e);
    }
  }

}


