package com.airbnb.aerosolve.core.transforms;

import com.airbnb.aerosolve.core.FeatureVector;
import com.airbnb.aerosolve.core.util.Util;
import com.typesafe.config.Config;

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

/**
 * Converts strings to either all lowercase or all uppercase
 * "field1" specifies the key of the feature
 * "convert_to_uppercase" converts strings to uppercase if true, otherwise converts to lowercase
 */
public class ConvertStringCaseTransform implements Transform {
  private String fieldName1;
  private boolean convertToUppercase;
  private String outputName;

  @Override
  public void configure(Config config, String key) {
    fieldName1 = config.getString(key + ".field1");
    convertToUppercase = config.getBoolean(key + ".convert_to_uppercase");
    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;
      String convertedString;
      if (convertToUppercase) {
        convertedString = rawString.toUpperCase();
      } else {
        convertedString = rawString.toLowerCase();
      }
      output.add(convertedString);
    }
  }
}

@@ -168,16 +168,43 @@
 onfig;%0A%0A
+import java.util.Iterator;%0A
 import j
@@ -447,16 +447,156 @@
 wercase%0A
+ * %22output%22 optionally specifies the key of the output feature, if it is not given the transform%0A * overwrites / replaces the input feature%0A
  */%0Apubl
@@ -939,24 +939,69 @@
 ppercase%22);%0A
+    if (config.hasPath(key + %22.output%22)) %7B%0A  
     outputNa
@@ -1040,16 +1040,66 @@
 tput%22);%0A
+    %7D else %7B%0A      outputName = fieldName1;%0A    %7D%0A
   %7D%0A%0A  @
@@ -1506,38 +1506,122 @@
 or (
-String rawString : feature1) %7B
+Iterator%3CString%3E iterator = feature1.iterator(); iterator.hasNext();) %7B%0A      String rawString = iterator.next();%0A
 %0A   
@@ -1637,25 +1637,25 @@
 wString 
-=
+!
 = null) 
 continue
@@ -1650,18 +1650,12 @@
 ll) 
-continue;%0A
+%7B%0A  
     
@@ -1682,24 +1682,26 @@
 ring;%0A      
+  
 if (convertT
@@ -1710,24 +1710,26 @@
 ppercase) %7B%0A
+  
         conv
@@ -1769,24 +1769,26 @@
 se();%0A      
+ 
+ 
 %7D else %7B%0A   
@@ -1780,24 +1780,26 @@
    %7D else %7B%0A
+  
         conv
@@ -1843,18 +1843,22 @@
 ;%0A      
+ 
+ 
 %7D%0A
+  
       ou
@@ -1884,16 +1884,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 ConvertStringCaseTransform