package com.airbnb.aerosolve.core.transforms;

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

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

import com.typesafe.config.Config;

/**
 * Takes the floats in fieldName1 and quantizes them into buckets. If the quantized float is
 * less than or equal to a maximum specified bucket value or greater than or equal to a minimum
 * specified bucket value, then the quantized float is stored as a string in a new string
 * feature output specified by stringOutputName. Otherwise, the original, unchanged float is
 * stored in a new float feature output specified by floatOutputName. The input float feature
 * remains unchanged.
 */
public class MoveFloatToStringAndFloatTransform implements Transform {
  private String fieldName1;
  private double bucket;
  private double maxBucket;
  private double minBucket;
  private String stringOutputName;
  private String floatOutputName;

  @Override
  public void configure(Config config, String key) {
    fieldName1 = config.getString(key + ".field1");
    bucket = config.getDouble(key + ".bucket");
    maxBucket = config.getDouble(key + ".max_bucket");
    minBucket = config.getDouble(key + ".min_bucket");
    stringOutputName = config.getString(key + ".string_output");
    floatOutputName = config.getString(key + ".float_output");
  }

  @Override
  public void doTransform(FeatureVector featureVector) {
    Map<String, Map<String, Double>> floatFeatures = featureVector.floatFeatures;

    if (floatFeatures == null || floatFeatures.isEmpty()) {
      return;
    }

    Map<String, Double> input = floatFeatures.get(fieldName1);

    if (input == null || input.isEmpty()) {
      return;
    }

    Util.optionallyCreateStringFeatures(featureVector);
    Map<String, Set<String>> stringFeatures = featureVector.getStringFeatures();
    Set<String> stringOutput = Util.getOrCreateStringFeature(stringOutputName, stringFeatures);

    Map<String, Double> floatOutput = Util.getOrCreateFloatFeature(floatOutputName, floatFeatures);

    for (Map.Entry<String, Double> inputEntry : input.entrySet()) {
      String inputFloatKey = inputEntry.getKey();
      Double inputFloatValue = inputEntry.getValue();

      Double inputFloatQuantized = LinearLogQuantizeTransform.quantize(inputFloatValue, bucket);

      if (inputFloatQuantized >= minBucket && inputFloatQuantized <= maxBucket) {
        String movedFloat = inputFloatKey + "=" + inputFloatQuantized;
        stringOutput.add(movedFloat);
      } else {
        floatOutput.put(inputFloatKey, inputFloatValue);
      }
    }
  }
}

@@ -133,16 +133,39 @@
 .Util;%0A%0A
+import java.util.List;%0A
 import j
@@ -178,16 +178,16 @@
 il.Map;%0A
-
 import j
@@ -268,22 +268,80 @@
  in 
-fieldName1 and
+the keys of fieldName1 (or if keys are not specified, all floats) and%0A *
  qua
@@ -891,16 +891,45 @@
 dName1;%0A
+  private List%3CString%3E keys;%0A
   privat
@@ -1184,24 +1184,121 @@
 %22.field1%22);%0A
+    if (config.hasPath(key + %22.keys%22)) %7B%0A      keys = config.getStringList(key + %22.keys%22);%0A    %7D%0A
     bucket =
@@ -2281,48 +2281,216 @@
     
-for (Map.Entry%3CString, Double%3E inputEntr
+if (keys != null) %7B%0A      for (String key : keys) %7B%0A        moveFloatToStringAndFloat(%0A          input, key, bucket, minBucket, maxBucket, stringOutput, floatOutput);%0A      %7D%0A    %7D else %7B%0A      for (String ke
 y : 
@@ -2499,12 +2499,10 @@
 put.
+k
 e
-ntr
 ySet
@@ -2517,51 +2517,411 @@
     
-String inputFloatKey = inputEntry.getKey();
+  moveFloatToStringAndFloat(%0A          input, key, bucket, minBucket, maxBucket, stringOutput, floatOutput);%0A      %7D%0A    %7D%0A  %7D%0A%0A  private static void moveFloatToStringAndFloat(%0A      Map%3CString, Double%3E input,%0A      String key,%0A      double bucket,%0A      double minBucket,%0A      double maxBucket,%0A      Set%3CString%3E stringOutput,%0A      Map%3CString, Double%3E floatOutput) %7B%0A    if (input.containsKey(key)) %7B
 %0A   
@@ -2957,23 +2957,16 @@
 nput
-Entry.getValue(
+.get(key
 );%0A%0A
@@ -3173,27 +3173,17 @@
 Float = 
-inputFloatK
+k
 ey + %22=%22
@@ -3287,19 +3287,9 @@
 put(
-inputFloatK
+k
 ey, 

Add keys to MoveFloatToStringAndFloatTransform