package com.airbnb.aerosolve.core.transforms;

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

import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class CapFloatFeatureTransform extends Transform {
  private String fieldName1;
  private List<String> keys;
  private double lowerBound;
  private double upperBound;
  private String output; // output family name, if not specified, output to fieldName1

  @Override
  public void configure(Config config, String key) {
    fieldName1 = config.getString(key + ".field1");
    keys = config.getStringList(key + ".keys");
    lowerBound = config.getDouble(key + ".lower_bound");
    upperBound = config.getDouble(key + ".upper_bound");
    if (config.hasPath(key + ".output")) {
      output = config.getString(key + ".output");
    } else {
      output = fieldName1;
    }
  }

  @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) {
      return;
    }
    Map<String, Double> feature2;
    if (output.equals(fieldName1)) {
      feature2 = feature1;
    } else if (floatFeatures.containsKey(output)) {
      feature2 = floatFeatures.get(output);
    } else {
      feature2 = new HashMap<>();
      floatFeatures.put(output, feature2);
    }

    for (String key : keys) {
      Double v = feature1.get(key);
      if (v != null) {
        feature2.put(key, Math.min(upperBound, Math.max(lowerBound, v)));
      }
    }
  }
}

@@ -88,16 +88,61 @@
 Vector;%0A
+import com.airbnb.aerosolve.core.util.Util;%0A%0A
 import c
@@ -438,16 +438,20 @@
 g output
+Name
 ; // out
@@ -841,16 +841,20 @@
   output
+Name
  = confi
@@ -908,16 +908,20 @@
   output
+Name
  = field
@@ -1298,144 +1298,28 @@
 ure2
-;%0A    if (output.equals(fieldName1)) %7B%0A      feature2 = feature1;%0A    %7D else if (floatFeatures.containsKey(output)) %7B%0A      feature2 = f
+ = Util.getOrCreateF
 loat
@@ -1329,119 +1329,36 @@
 ture
-s.get
 (output
-);%0A    %7D else %7B%0A      feature2 = new HashMap%3C%3E();%0A      floatFeatures.put(output, feature2);%0A    %7D%0A
+Name, floatFeatures);
 %0A   

simplify the logic