package com.airbnb.aerosolve.core.transforms;

import com.airbnb.aerosolve.core.FeatureVector;
import com.airbnb.aerosolve.core.util.Util;

import java.util.List;
import java.util.Map;
import java.util.Set;

import com.typesafe.config.Config;
import com.typesafe.config.ConfigObject;

/**
 * Replaces all substrings that match a given regex with a replacement string
 * "field1" specifies the key of the feature
 * "replacements" specifies a list of pairs (or maps) of regexes and corresponding replacements
 * Replacements are performed in the same order as specified in the list of pairs
 * "replacement" specifies the replacement string
 */
public class ReplaceAllStringsTransform implements Transform {
  private String fieldName1;
  private List<? extends ConfigObject> replacements;
  private String outputName;

  @Override
  public void configure(Config config, String key) {
    fieldName1 = config.getString(key + ".field1");
    replacements = config.getObjectList(key + ".replacements");
    outputName = config.getString(key + ".output");
  }

  @Override
  public void doTransform(FeatureVector featureVector) {
    Map<String, Set<String>> stringFeatures = featureVector.getStringFeatures();
    if (stringFeatures == null) {
      return;
    }

    Set<String> feature1 = stringFeatures.get(fieldName1);
    if (feature1 == null) {
      return;
    }

    Set<String> output = Util.getOrCreateStringFeature(outputName, stringFeatures);

    for (String rawString : feature1) {
      if (rawString == null) continue;
      for (ConfigObject replacementCO : replacements) {
        Map<String, Object> replacementMap = replacementCO.unwrapped();
        for (Map.Entry<String, Object> replacementEntry : replacementMap.entrySet()) {
          String regex = replacementEntry.getKey();
          String replacement = (String) replacementEntry.getValue();
          rawString = rawString.replaceAll(regex, replacement);
        }
      }
      output.add(rawString);
    }
  }
}

@@ -133,16 +133,43 @@
 .Util;%0A%0A
+import java.util.Iterator;%0A
 import j
@@ -1016,24 +1016,69 @@
 acements%22);%0A
+    if (config.hasPath(key + %22.output%22)) %7B%0A  
     outputNa
@@ -1117,16 +1117,66 @@
 tput%22);%0A
+    %7D else %7B%0A      outputName = fieldName1;%0A    %7D%0A
   %7D%0A%0A  @
@@ -1583,38 +1583,122 @@
 or (
-String rawString : feature1) %7B
+Iterator%3CString%3E iterator = feature1.iterator(); iterator.hasNext();) %7B%0A      String rawString = iterator.next();%0A
 %0A   
@@ -1714,25 +1714,25 @@
 wString 
-=
+!
 = null) 
 continue
@@ -1727,18 +1727,12 @@
 ll) 
-continue;%0A
+%7B%0A  
     
@@ -1787,24 +1787,26 @@
 ) %7B%0A        
+ 
+ 
 Map%3CString, 
@@ -1853,24 +1853,27 @@
 nwrapped();%0A
+%0A  
         for 
@@ -1949,32 +1949,34 @@
 ()) %7B%0A          
+ 
+ 
 String regex = r
@@ -1997,24 +1997,26 @@
 y.getKey();%0A
+  
           St
@@ -2068,24 +2068,26 @@
 getValue();%0A
+  
           ra
@@ -2150,18 +2150,22 @@
     
-%7D%0A
+  %7D%0A  
       %7D%0A
     
@@ -2156,24 +2156,26 @@
 %7D%0A        %7D%0A
+  
       output
@@ -2191,16 +2191,185 @@
 tring);%0A
+      %7D%0A%0A      // Check reference equality to determine whether the output should overwrite the input%0A      if (output == feature1) %7B%0A        iterator.remove();%0A      %7D%0A
     %7D%0A  

Make output optional for ReplaceAllStringsTransform