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.Set;
import java.util.HashSet;
import java.util.Map;
import java.util.Map.Entry;
import java.util.List;

/**
 * Moves named fields from one family to another.
 */
public class MoveFloatToStringTransform implements Transform {
  private String fieldName1;
  private double bucket;
  private String outputName;
  private List<String> keys;
  private double cap;

  @Override
  public void configure(Config config, String key) {
    fieldName1 = config.getString(key + ".field1");
    bucket = config.getDouble(key + ".bucket");
    outputName = config.getString(key + ".output");
    keys = config.getStringList(key + ".keys");
    try {
      cap = config.getDouble(key + ".cap");
    } catch (Exception e) {
      cap = 1e10;
    }
  }

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

    if (floatFeatures == null) {
      return;
    }
    Map<String, Double> feature1 = floatFeatures.get(fieldName1);
    if (feature1 == null || feature1.isEmpty()) {
      return;
    }

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

    for (String key : keys) {
      if (feature1.containsKey(key)) {
        Double dbl = feature1.get(key);
        if (dbl > cap) {
          dbl = cap;
        }

        Double quantized = LinearLogQuantizeTransform.quantize(dbl, bucket);
        output.add(key + '=' + quantized);
        feature1.remove(key);
      }
    }
  }
}

@@ -342,16 +342,88 @@
 another.
+ If keys are not specified, all keys are moved%0A * from the float family.
 %0A */%0Apub
@@ -830,24 +830,67 @@
 %22.output%22);%0A
+    if (config.hasPath(key + %22.keys%22)) %7B%0A  
     keys = c
@@ -929,19 +929,55 @@
 %22);%0A    
-try
+%7D%0A    if (config.hasPath(key + %22.cap%22))
  %7B%0A     
@@ -1025,27 +1025,12 @@
   %7D 
-catch (Exception e)
+else
  %7B%0A 
@@ -1622,24 +1622,50 @@
 Features);%0A%0A
+    if (keys != null) %7B%0A  
     for (Str
@@ -1680,24 +1680,370 @@
 : keys) %7B%0A  
+      moveFloat(feature1, output, key, cap, bucket);%0A      %7D%0A    %7D else %7B%0A      for (String key : feature1.keySet()) %7B%0A        moveFloat(feature1, output, key, cap, bucket);%0A      %7D%0A    %7D%0A  %7D%0A%0A  public static void moveFloat(%0A      Map%3CString, Double%3E feature1,%0A      Set%3CString%3E output,%0A      String key,%0A      double cap,%0A      double bucket) %7B%0A
     if (feat
@@ -2063,26 +2063,24 @@
 Key(key)) %7B%0A
-  
       Double
@@ -2101,26 +2101,24 @@
 1.get(key);%0A
-  
       if (db
@@ -2136,18 +2136,16 @@
         
-  
 dbl = ca
@@ -2153,23 +2153,19 @@
 ;%0A      
-  
 %7D%0A%0A
-  
       Do
@@ -2237,18 +2237,16 @@
 ;%0A      
-  
 output.a
@@ -2278,18 +2278,16 @@
 ;%0A      
-  
 feature1
@@ -2300,24 +2300,16 @@
 e(key);%0A
-      %7D%0A
     %7D%0A  

Make keys optional for MoveFloatToStringTransform