package com.airbnb.aerosolve.core.transforms;

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

import java.io.Serializable;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

public class Transformer implements Serializable {

  private static final long serialVersionUID = 1569952057032186608L;
  // The transforms to be applied to the context, item and combined
  // (context | item) respectively.
  private final Transform contextTransform;
  private final Transform itemTransform;
  private final Transform combinedTransform;

  public Transformer(Config config, String key) {
    // Configures the model transforms.
    // context_transform : name of ListTransform to apply to context
    // item_transform : name of ListTransform to apply to each item
    // combined_transform : name of ListTransform to apply to each (item context) pair
    String contextTransformName = config.getString(key + ".context_transform");
    contextTransform = TransformFactory.createTransform(config, contextTransformName);
    String itemTransformName = config.getString(key + ".item_transform");
    itemTransform = TransformFactory.createTransform(config, itemTransformName);
    String combinedTransformName = config.getString(key + ".combined_transform");
    combinedTransform = TransformFactory.createTransform(config, combinedTransformName);
  }

  // Helper functions for transforming context, items or combined feature vectors.
  public void transformContext(FeatureVector context) {
    if (contextTransform != null && context != null) {
      contextTransform.doTransform(context);
    }
  }

  public void transformItem(FeatureVector item) {
    if (itemTransform != null && item != null) {
      itemTransform.doTransform(item);
    }
  }

  public void transformItems(List<FeatureVector> items) {
    if (items != null) {
      for (FeatureVector item : items) {
        transformItem(item);
      }
    }
  }

  public void transformCombined(FeatureVector combined) {
    if (combinedTransform != null && combined != null) {
      combinedTransform.doTransform(combined);
    }
  }

  // In place apply all the transforms to the context and items
  // and apply the combined transform to items.
  public void combineContextAndItems(Example examples) {
    transformContext(examples.context);
    transformItems(examples.example);
    addContextToItemsAndTransform(examples);
  }

  // Adds the context to items and applies the combined transform
  public void addContextToItemsAndTransform(Example examples) {
    Map<String, Set<String>> contextStringFeatures = null;
    if (examples.context != null &&
        examples.context.stringFeatures != null) {
      contextStringFeatures = examples.context.getStringFeatures();
    }
    for (FeatureVector item : examples.example) {
      addContextToItemAndTransform(contextStringFeatures, item);
    }
  }

  public void addContextToItemAndTransform(Map<String, Set<String>> contextStringFeatures,
                                           FeatureVector item) {
    Map<String, Set<String>> itemStringFeatures = item.getStringFeatures();
    if (item.getStringFeatures() == null) {
      item.setStringFeatures(contextStringFeatures);
    } else if (contextStringFeatures != null) {
      itemStringFeatures.putAll(contextStringFeatures);
    }
    transformCombined(item);
  }
}

@@ -2689,16 +2689,141 @@
 = null;%0A
+    Map%3CString, Map%3CString, Double%3E%3E contextFloatFeatures = null;%0A    Map%3CString, List%3CDouble%3E%3E contextDenseFeatures = null;%0A
     if (
@@ -2850,20 +2850,22 @@
 null
- &&
+) %7B
 %0A      
-  
+if (
 exam
@@ -2905,24 +2905,26 @@
 ll) %7B%0A      
+  
 contextStrin
@@ -2969,24 +2969,288 @@
 Features();%0A
+      %7D%0A      if (examples.context.floatFeatures != null) %7B%0A        contextFloatFeatures = examples.context.getFloatFeatures();%0A      %7D%0A      if (examples.context.denseFeatures != null) %7B%0A        contextDenseFeatures = examples.context.getDenseFeatures();%0A      %7D%0A
     %7D%0A    fo
@@ -3324,24 +3324,35 @@
 ndTransform(
+%0A          
 contextStrin
@@ -3361,16 +3361,60 @@
 eatures,
+ contextFloatFeatures, contextDenseFeatures,
  item);%0A
@@ -3562,296 +3562,1044 @@
     
-FeatureVector item) %7B%0A    Map%3CString, Set%3CString%3E%3E itemStringFeatures = item.getStringFeatures();%0A    if (item.getStringFeatures() == null) %7B%0A      item.setStringFeatures(contextStringFeatures);%0A    %7D else if (contextStringFeatures != null) %7B%0A      itemStringFeatures.putAll(contextString
+Map%3CString, Map%3CString, Double%3E%3E contextFloatFeatures,%0A                                           Map%3CString, List%3CDouble%3E%3E contextDenseFeatures,%0A                                           FeatureVector item) %7B%0A    Map%3CString, Set%3CString%3E%3E itemStringFeatures = item.getStringFeatures();%0A    Map%3CString, Map%3CString, Double%3E%3E itemFloatFeatures = item.getFloatFeatures();%0A    Map%3CString, List%3CDouble%3E%3E itemDenseFeatures = item.getDenseFeatures();%0A    if (item.getStringFeatures() == null) %7B%0A      item.setStringFeatures(contextStringFeatures);%0A    %7D else if (contextStringFeatures != null) %7B%0A      itemStringFeatures.putAll(contextStringFeatures);%0A    %7D%0A    if (item.getFloatFeatures() == null) %7B%0A      item.setFloatFeatures(contextFloatFeatures);%0A    %7D else if (contextFloatFeatures != null) %7B%0A      itemFloatFeatures.putAll(contextFloatFeatures);%0A    %7D%0A    if (item.getDenseFeatures() == null) %7B%0A      item.setDenseFeatures(contextDenseFeatures);%0A    %7D else if (contextDenseFeatures != null) %7B%0A      itemDenseFeatures.putAll(contextDense
 Feat

Copy all features from context to example