package com.acumenvelocity.ath.solr.tm;
   
   import java.util.ArrayList;
   import java.util.Collections;
   import java.util.HashMap;
   import java.util.List;
   import java.util.Map;
   
   import org.apache.solr.client.solrj.response.QueryResponse;
  import org.apache.solr.common.SolrDocument;
  import org.apache.solr.common.SolrDocumentList;
  
  import com.acumenvelocity.ath.common.Const;
  import com.acumenvelocity.ath.common.ConversionUtil;
  import com.acumenvelocity.ath.common.JacksonUtil;
  import com.acumenvelocity.ath.common.Log;
  import com.acumenvelocity.ath.common.OkapiUtil;
  import com.acumenvelocity.ath.common.SolrUtil;
  import com.acumenvelocity.ath.common.exception.AthRuntimeException;
  import com.acumenvelocity.ath.model.x.LayeredTextX;
  import com.acumenvelocity.ath.solr.AthIndex;
  import com.acumenvelocity.ath.solr.BigramAnalyzer;
  import com.acumenvelocity.ath.solr.Solr;
  
  import net.sf.okapi.common.IParameters;
  import net.sf.okapi.common.query.MatchType;
  import net.sf.okapi.common.query.QueryResult;
  import net.sf.okapi.common.resource.TextFragment;
  import net.sf.okapi.lib.translation.BaseConnector;
  import net.sf.okapi.lib.translation.ITMQuery;
  
  public class SolrTmConnector extends BaseConnector implements ITMQuery {
    private Parameters params;
    private List<QueryResult> results;
    private int maxHits = 3;
    private int threshold = 75;
    private int current = -1;
    private BigramAnalyzer bigramAnalyzer = new BigramAnalyzer();
  
    // The fields to be returned
    public static final String[] EXACT_RETURN_FIELDS = { Const.ATH_PROP_SOURCE_JSON,
        Const.ATH_PROP_TARGET_JSON, Const.ATH_PROP_SOLR_ID, Const.ATH_PROP_TM_SEG_ID,
        Const.ATH_PROP_SOLR_SCORE };
  
    public static final String[] FUZZY_RETURN_FIELDS = { Const.ATH_PROP_SOURCE,
        Const.ATH_PROP_SOURCE_JSON, Const.ATH_PROP_TARGET, Const.ATH_PROP_TARGET_JSON,
        Const.ATH_PROP_SOLR_ID, Const.ATH_PROP_TM_SEG_ID, Const.ATH_PROP_SOLR_SCORE };
  
    private static final int MAX_ROWS = 500;
  
    public SolrTmConnector() {
      super();
      setParameters(new Parameters());
      setWeight(1);
    }
  
    @Override
    public String getName() {
      return "Solr TM";
    }
  
    @Override
    public String getSettingsDisplay() {
      return Log.format(
          "\nTM ID: {}\n"
              + "Solr core name: {}\n"
              + "Search in the source: {}\n"
              + "Penalize TM matches with different tags in source: {}, in target: {}\n",
  
          net.sf.okapi.common.Util.isEmpty(params.getTmId()) ? "<To be specified>"
              : params.getTmId(),
  
          net.sf.okapi.common.Util.isEmpty(params.getSolrCoreName()) ? "<To be specified>"
              : params.getSolrCoreName(),
  
          params.getSearchInSource() ? "Yes" : "No",
          params.getPenalizeSourceTagsDifference() ? "Yes" : "No",
          params.getPenalizeTargetTagsDifference() ? "Yes" : "No");
    }
  
    @Override
    public boolean hasNext() {
      if (results == null) {
        return false;
      }
  
      if (current >= results.size()) {
        current = -1;
      }
  
      return (current > -1);
    }
  
    @Override
    public QueryResult next() {
      if (results == null) {
        return null;
      }
  
     if ((current > -1) && (current < results.size())) {
       current++;
       return results.get(current - 1);
     }
 
     current = -1;
     return null;
   }
 
   @Override
   public void open() {
     // Nothing
   }
 
   @Override
   public void close() {
     // Nothing
   }
 
   @Override
   public int query(String plainText) {
     return query(new TextFragment(plainText));
   }
 
   @Override
   public int query(TextFragment queryTf) {
     // P.s(Const.PROFILER_ZONE_TM_CONNECTOR, "Solr connector");
     try {
       // P.s(15, "Prepare search");
       results = new ArrayList<>();
       current = -1;
 
       LayeredTextX queryLt = ConversionUtil.toLayeredText(queryTf, super.srcLoc);
       String queryText = queryTf.getText();
       Log.trace(this.getClass(), "--- query: {}", queryText);
       String queryCodes = queryTf.getCodes().toString();
 
       Map<String, Object> queryParams = new HashMap<>();
       queryParams.put(Solr.INCLUDE_SCORE, true);
       // P.e(15);
 
       try {
         // P.s(17, "Both exact & fuzzy");
         // P.s(Const.PROFILER_ZONE_TM_SOLR_EXACT, "Solr exact");
 
         // Query Solr for an exact match of the text with tags
         queryParams.put(Solr.FIELDS, EXACT_RETURN_FIELDS);
 
         // XXX We need only one exact match, only single target TMs are supported
         queryParams.put(Solr.ROWS, 1);
         String exactId = null;
 
         // Ignore those with hidden=true (for MT cache use only)
         // P.s(18, "Solr exact query");
         QueryResponse response = AthIndex.getMany(params.getSolrCoreName(),
             Log.format(
                 params.getSearchInSource()
                     ? "tmId:\"{}\" AND srcLang:\"{}\" AND trgLang:\"{}\" AND sourceWithCodes:\"{}\""
                     : "tmId:\"{}\" AND srcLang:\"{}\" AND trgLang:\"{}\" AND targetWithCodes:\"{}\"",
                 params.getTmId(),
                 super.srcCode,
                 super.trgCode,
                 SolrUtil.normalizeQuery(queryLt.getTextWithCodes())),
             queryParams, QueryResponse.class);
         // P.e(18);
 
         SolrDocumentList docList = response.getResults();
         if (docList.getNumFound() > 0) {
           // Exact match
           SolrDocument doc = docList.get(0);
 
           String srcJson = doc._getStr(Const.ATH_PROP_SOURCE_JSON, null);
           String trgJson = doc._getStr(Const.ATH_PROP_TARGET_JSON, null);
 
           if (srcJson != null && trgJson != null) {
             LayeredTextX srcSc = JacksonUtil.fromJson(srcJson, LayeredTextX.class);
 
             LayeredTextX trgSc = JacksonUtil.fromJson(trgJson, LayeredTextX.class);
 
             if (srcSc == null || trgSc == null) {
               Log.warn(this.getClass(), "Broken Solr document, id='{}'", doc._getStr("id", null));
 
             } else {
               TextFragment tmSrcTf = ConversionUtil.toTextFragment(srcSc);
               TextFragment tmTrgTf = ConversionUtil.toTextFragment(trgSc);
 
               QueryResult qr = new QueryResult();
               results.add(qr);
 
               qr.weight = super.getWeight();
               qr.setFuzzyScore(100);
               qr.entryId = doc._getStr(Const.ATH_PROP_SOLR_ID, null);
               qr.source = tmSrcTf;
               qr.target = tmTrgTf;
               qr.matchType = MatchType.EXACT;
               qr.origin = getName();
 
               exactId = qr.entryId;
 
               Log.trace(this.getClass(), "\n--- {}%\n--seg: {}\n---tm: {}", qr.getFuzzyScore(),
                   queryLt.getText(),
                   tmSrcTf.getText());
             }
           }
         }
 
         // P.e(Const.PROFILER_ZONE_TM_SOLR_EXACT);
         // P.s(Const.PROFILER_ZONE_TM_SOLR_FUZZY, "Solr fuzzy");
 
         if (maxHits > results.size()) {
           // Query Solr fuzzily for the text without tags
           // XXX We return only the highest score fuzzy match
           queryParams.put(Solr.FIELDS, FUZZY_RETURN_FIELDS);
 
           // // +1 for a duplicate exact match among fuzzy ones
           // int numFuzzies = maxHits - results.size();
           // if (results.size() > 0) {
           // numFuzzies++;
           // }
 
           // queryParams.put(Solr.ROWS, numFuzzies);
           queryParams.put(Solr.ROWS, MAX_ROWS);
           // Ignore those with hidden=true (for MT cache use only)
           // P.s(19, "Solr fuzzy query");
           response = AthIndex.getMany(params.getSolrCoreName(),
               Log.format(
                   params.getSearchInSource()
                       ? "tmId:\"{}\" AND srcLang:\"{}\" AND trgLang:\"{}\" AND source:\"{}\""
                       : "tmId:\"{}\" AND srcLang:\"{}\" AND trgLang:\"{}\" AND target:\"{}\"",
                   params.getTmId(),
                   super.srcCode,
                   super.trgCode,
                   SolrUtil.normalizeQuery(queryLt.getText())),
               queryParams, QueryResponse.class);
           // P.e(19);
 
           // P.c(16, "QTime", response.getQTime(), queryText);
 
           docList = response.getResults();
           if (docList.getNumFound() > 0) {
             if (docList.getNumFound() > MAX_ROWS) {
               Log.warn(getClass(), "Solr fuzzy query returned too many documents: {}",
                   docList.getNumFound());
             }
             // int maxScore = 0;
             // SolrDocument maxScoreDoc = null;
 
             for (SolrDocument doc : docList) {
               String docId = doc._getStr(Const.ATH_PROP_SOLR_ID, null);
 
               // Don't add to results what has been matched exactly before
               if (docId != null && docId.equalsIgnoreCase(exactId)) {
                 continue;
               }
 
               String tmText = (String) doc.getFieldValue(
                   params.getSearchInSource() ? Const.ATH_PROP_SOURCE : Const.ATH_PROP_TARGET);
 
               if (!net.sf.okapi.common.Util.isEmpty(tmText)) {
                 int score = Math.round(OkapiUtil.calculateNgramDiceCoefficient(
                     tmText, queryText, bigramAnalyzer));
 
                 // // XXX If there are several matches with the same score, the first one is
                 // kept
                 // if (maxScore < score) {
                 // maxScore = score;
                 // maxScoreDoc = doc;
                 // }
 
                 String srcJson = doc._getStr(Const.ATH_PROP_SOURCE_JSON, null);
                 String trgJson = doc._getStr(Const.ATH_PROP_TARGET_JSON, null);
 
                 if (srcJson != null && trgJson != null) {
                   LayeredTextX srcSc = JacksonUtil.fromJson(srcJson, LayeredTextX.class);
 
                   LayeredTextX trgSc = JacksonUtil.fromJson(trgJson, LayeredTextX.class);
 
                   if (srcSc == null || trgSc == null) {
                     Log.warn(this.getClass(), "Broken Solr document, id='{}'",
                         doc._getStr(Const.ATH_PROP_SOLR_ID, null));
 
                   } else {
                     TextFragment tmSrcTf = ConversionUtil.toTextFragment(srcSc);
                     TextFragment tmTrgTf = ConversionUtil.toTextFragment(trgSc);
 
                     // Apply penalties
                     // XXX If another match has equal codes with the query, but a lower text
                     // matching score, it is still a lower match than the one with the text closer
                     // to the query, but probably worse codes. We give a higher priority to the
                     // text, not the codes.
                     if (params.getPenalizeSourceTagsDifference()) {
                       if (!queryCodes.equals(tmSrcTf.getCodes().toString())) {
                         score--; // 1% penalty
                       }
                     }
 
                     if (params.getPenalizeTargetTagsDifference()) {
                       if (!queryCodes.equals(tmTrgTf.getCodes().toString())) {
                         score--; // 1% penalty
                       }
                     }
 
                     if (score < threshold) {
                       Log.trace(this.getClass(), "\n--- {}%\n--seg: {}\n---tm: {}", score,
                           queryLt.getText(),
                           tmSrcTf.getText());
 
                       // Drop the hit
                       continue;
                     }
 
                     QueryResult qr = new QueryResult();
                     results.add(qr);
 
                     qr.weight = super.getWeight();
 
                     if (SolrUtil.checkTmFuzzyScore(score)) {
                       qr.setFuzzyScore(score);
                     }
 
                     qr.entryId = doc._getStr(Const.ATH_PROP_SOLR_ID, null);
                     qr.source = tmSrcTf;
                     qr.target = tmTrgTf;
                     qr.matchType = score >= 100 ? MatchType.EXACT : MatchType.FUZZY;
                     qr.origin = getName();
 
                     Log.trace(this.getClass(), "\n--- {}%\n--seg: {}\n---tm: {}",
                         qr.getFuzzyScore(),
                         queryLt.getText(),
                         tmSrcTf.getText());
                   }
                 }
               }
             }
           }
 
           // Sort QueryResults on MatchType, Score, and Source String
           Collections.sort(results);
         }
 
         // P.e(Const.PROFILER_ZONE_TM_SOLR_FUZZY);
         // P.e(17);
 
       } catch (Exception e) {
         AthRuntimeException.logAndThrow(SolrTmConnector.class, e);
       }
 
       if (results.size() > 0) {
         current = 0;
       }
 
       return results.size();
 
     } finally {
       // P.e(Const.PROFILER_ZONE_TM_CONNECTOR);
     }
   }
 
   @Override
   public void setMaximumHits(int maxHits) {
     if (maxHits < 1) {
       this.maxHits = 1;
 
     } else {
       this.maxHits = maxHits;
     }
   }
 
   @Override
   public int getMaximumHits() {
     return maxHits;
   }
 
   @Override
   public int getThreshold() {
     return threshold;
   }
 
   @Override
   public void setThreshold(int threshold) {
     this.threshold = threshold;
   }
 
   @Override
   public IParameters getParameters() {
     return params;
   }
 
   @Override
   public void setParameters(IParameters params) {
     this.params = (Parameters) params;
   }
 }