package org.oscii.concordance;

import edu.stanford.nlp.mt.tm.DynamicTranslationModel;
import edu.stanford.nlp.mt.tm.SampledRule;
import edu.stanford.nlp.mt.util.ParallelSuffixArray;
import gnu.trove.map.hash.THashMap;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.oscii.lex.Expression;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Stream;

import static java.util.Collections.emptyList;
import static java.util.Collections.emptyMap;
import static java.util.stream.Collectors.counting;
import static java.util.stream.Collectors.groupingBy;

/**
 * Corpus backed by a suffix array.
 */
public class SuffixArrayCorpus extends AlignedCorpus {
    // source language -> target language -> suffix array
    Map<String, Map<String, ParallelSuffixArray>> suffixes = new HashMap<>();
    private int maxSamples = 100;
    private int maxTargetPhrase = 5;

    private final static Logger log = LogManager.getLogger(SuffixArrayCorpus.class);

    @Override
    public void read(String path, String sourceLanguage, String targetLanguage, int max) throws IOException {
        log.info("Reading sentences: " + sourceLanguage + "-" + targetLanguage);
        ParallelFiles paths = paths(path, sourceLanguage, targetLanguage);
        ParallelSuffixArray suffixArray = new ParallelSuffixArray(
                paths.sourceSentences.toString(),
                paths.targetSentences.toString(),
                paths.alignments.toString(),
                10000);
        // TODO(spenceg) Fix build
//        suffixArray.createRuleCaches(maxSamples, 1000);
        Map<String, ParallelSuffixArray> bySource = suffixes.get(sourceLanguage);
        if (bySource == null) {
            bySource = new THashMap<>();
            suffixes.put(sourceLanguage, bySource);
        }
        if (bySource.containsKey(targetLanguage)) {
            throw new RuntimeException("Multiple corpora for a language pair: "
                    + sourceLanguage + ", " + targetLanguage);
        }
        bySource.put(targetLanguage, suffixArray);
    }

    @Override
    public Function<Expression, Double> translationFrequencies(Expression source) {
        Map<String, ParallelSuffixArray> arrays = suffixes.get(source.language);
        if (arrays == null) {
            return AlignedCorpus::zeroFrequency;
        }

        Map<String, Map<String, Long>> counts = new THashMap<>();
        arrays.entrySet().forEach(kv -> counts.put(kv.getKey(), countAll(source.text, kv.getValue())));
        return normalizeByLanguage(counts);
    }

    private Map<String, Long> countAll(String text, ParallelSuffixArray suffixArray) {
        // Generate samples
        String[] words = text.trim().split("\\s+");
        int[] phrase = new int[words.length];
        for (int i = 0; i < phrase.length; i++) {
            phrase[i] = suffixArray.getVocabulary().indexOf(words[i]);
            if (phrase[i] == -1) {
                return emptyMap();
            }
        }
        List<ParallelSuffixArray.QueryResult> samples = suffixArray.sample(phrase, true, maxSamples).samples;

        // Count translations
        // TODO(spenceg) Comment out to fix build
//        Stream<SampledRule> rules = samples.stream().flatMap(
//                s -> DynamicTranslationModel.extractRules(s, words.length, maxTargetPhrase).stream());
//        return rules.collect(groupingBy(rule -> targetOf(rule, suffixArray), counting()));
        return null;
    }

    private String targetOf(SampledRule rule, ParallelSuffixArray suffixArray) {
        String[] words = new String[rule.tgt.length];
        for (int i = 0; i < rule.tgt.length; i++) {
            words[i] = suffixArray.getVocabulary().get(rule.tgt[i]);
        }
        return String.join(" ", words);
    }

    @Override
    public List<AlignedSentence> examples(String query, String source, String target, int max) {
        return emptyList();
    }
}

@@ -1363,32 +1363,69 @@
 argetLanguage);%0A
+        // TODO(spenceg) Fix build%0A//
         Parallel
@@ -1475,16 +1475,18 @@
 xArray(%0A
+//
         
@@ -1519,32 +1519,34 @@
 ces.toString(),%0A
+//
                 
@@ -1571,32 +1571,34 @@
 ces.toString(),%0A
+//
                 
@@ -1626,16 +1626,18 @@
 ring(),%0A
+//
         
@@ -1656,43 +1656,8 @@
 0);%0A
-        // TODO(spenceg) Fix build%0A
 //  
@@ -1710,16 +1710,18 @@
  1000);%0A
+//
         
@@ -1786,32 +1786,34 @@
 ourceLanguage);%0A
+//
         if (bySo
@@ -1820,32 +1820,34 @@
 urce == null) %7B%0A
+//
             bySo
@@ -1863,32 +1863,34 @@
 w THashMap%3C%3E();%0A
+//
             suff
@@ -1921,24 +1921,26 @@
  bySource);%0A
+//
         %7D%0A  
@@ -1929,32 +1929,34 @@
 e);%0A//        %7D%0A
+//
         if (bySo
@@ -1987,24 +1987,26 @@
 anguage)) %7B%0A
+//
             
@@ -2073,16 +2073,18 @@
 pair: %22%0A
+//
         
@@ -2130,32 +2130,34 @@
 argetLanguage);%0A
+//
         %7D%0A      
@@ -2142,32 +2142,34 @@
 e);%0A//        %7D%0A
+//
         bySource
@@ -3165,19 +3165,20 @@
 ray.
-QueryResult
+SentencePair
 %3E sa
@@ -3215,14 +3215,8 @@
 ase,
- true,
  max

Fix build