package ai.h2o.targetencoding;

import hex.ModelBuilder;
import hex.ModelCategory;
import water.DKV;
import water.Scope;
import water.fvec.Frame;
import water.fvec.Vec;
import water.util.IcedHashMapGeneric;
import water.util.Log;

import java.util.*;

public class TargetEncoderBuilder extends ModelBuilder<TargetEncoderModel, TargetEncoderModel.TargetEncoderParameters, TargetEncoderModel.TargetEncoderOutput> {

  public TargetEncoderModel getTargetEncoderModel() {
    assert _targetEncoderModel != null : "Training phase of the TargetEncoderBuilder did not take place yet. TargetEncoderModel is not available.";
    return _targetEncoderModel;
  }

  private TargetEncoderModel _targetEncoderModel;
  
  public TargetEncoderBuilder(TargetEncoderModel.TargetEncoderParameters parms) {
    super(parms);
    super.init(false);
  }

  public TargetEncoderBuilder(final boolean startupOnce) {
    super(new TargetEncoderModel.TargetEncoderParameters(), startupOnce);
  }

  @Override
  public void init(boolean expensive) {
    super.init(expensive);
  }

  private class TargetEncoderDriver extends Driver {
    @Override
    public void computeImpl() {
      // In case ModelBuilder is registered with default model parameters ( i.e. without training frame etc.) we want to init it at this stage as well
      if(_startUpOnceModelBuilder) init(false);

      final int numColsRemoved = hasFoldCol() ? 2 : 1; // Response is always at the last index, fold column is on the index before.

      final String[] encodedColumns = Arrays.copyOf(train().names(), train().names().length - numColsRemoved);
      TargetEncoder tec = new TargetEncoder(encodedColumns);

      Scope.untrack(train().keys());

      IcedHashMapGeneric<String, Frame> _targetEncodingMap = tec.prepareEncodingMap(train(), _parms._response_column, _parms._fold_column);

      // Mean could be computed from any encoding map as response column is shared
      double priorMean = tec.calculatePriorMean(_targetEncodingMap.entrySet().iterator().next().getValue());

      for(Map.Entry<String, Frame> entry: _targetEncodingMap.entrySet()) {
        Frame frameWithEncodingMap = entry.getValue();
        Scope.untrack(frameWithEncodingMap.keys());
      }

      disableIgnoreConstColsFeature();
      TargetEncoderModel.TargetEncoderOutput output = new TargetEncoderModel.TargetEncoderOutput(TargetEncoderBuilder.this, _targetEncodingMap, priorMean);
      _targetEncoderModel = new TargetEncoderModel(_result, _parms, output, tec);
      DKV.put(_targetEncoderModel);
    }

    private void disableIgnoreConstColsFeature() {
      _parms._ignore_const_cols = false;
      Log.info("We don't want to ignore any columns during target encoding transformation therefore `_ignore_const_cols` parameter was set to `false`");
    }
  }

  @Override
  protected void ignoreInvalidColumns(int npredictors, boolean expensive) {
    new FilterCols(npredictors){
      @Override
      protected boolean filter(Vec v) {
        return !v.isCategorical();
      }
    }.doIt(train(), "Removing non-categorical columns found in the list of encoded columns.", expensive);
  }

  /**
   * Never do traditional cross-validation for Target Encoder Model. The {@link TargetEncoder} class handles
   * fold column on it's own.
   *
   * @return Always false
   */
  @Override
  public boolean nFoldCV() {
    return false;
  }

  @Override
  protected Driver trainModelImpl() {
    // We can use Model.Parameters to configure Target Encoder
    return new TargetEncoderDriver();
  }
  
  @Override
  public ModelCategory[] can_build() {
    return new ModelCategory[]{ ModelCategory.TargetEncoder};
  }

  @Override
  public boolean isSupervised() {
    return true;
  }

  @Override
  public BuilderVisibility builderVisibility() {
    return BuilderVisibility.Stable;
  }

  @Override
  public boolean haveMojo() {
    return true;
  }

  @Override
  public String getName() {
    return "targetencoder";
  }
}

@@ -969,92 +969,8 @@
  %7D%0A%0A
-  @Override%0A  public void init(boolean expensive) %7B%0A    super.init(expensive);%0A  %7D%0A%0A
   pr
@@ -1074,202 +1074,71 @@
     
-// In case ModelBuilder is registered with default model parameters ( i.e. without training frame etc.) we want to init it at this stage as well%0A      if(_startUpOnceModelBuilder) init(fals
+_targetEncoderModel = null;%0A      try %7B%0A        init(tru
 e);%0A%0A
+  
     
@@ -1265,17 +1265,18 @@
 before.%0A
-%0A
+  
       fi
@@ -1378,16 +1378,18 @@
 moved);%0A
+  
       Ta
@@ -1438,24 +1438,417 @@
 dColumns);%0A%0A
+        TargetEncoderModel.TargetEncoderOutput emptyOutput =%0A                new TargetEncoderModel.TargetEncoderOutput(TargetEncoderBuilder.this, new IcedHashMapGeneric%3C%3E(), Double.NaN);%0A        TargetEncoderModel model = new TargetEncoderModel(dest(), _parms, emptyOutput, tec);%0A        _targetEncoderModel = model.delete_and_lock(_job); // and clear & write-lock it (smashing any prior)%0A%0A  
       Scope.
@@ -1869,24 +1869,26 @@
 ).keys());%0A%0A
+  
       IcedHa
@@ -2022,16 +2022,18 @@
 %0A%0A      
+  
 // Mean 
@@ -2101,16 +2101,18 @@
  shared%0A
+  
       do
@@ -2219,16 +2219,18 @@
 %0A%0A      
+  
 for(Map.
@@ -2286,24 +2286,26 @@
 trySet()) %7B%0A
+  
         Fram
@@ -2347,32 +2347,34 @@
 alue();%0A        
+  
 Scope.untrack(fr
@@ -2403,27 +2403,31 @@
 s());%0A      
+ 
+ 
 %7D%0A%0A
+  
       disabl
@@ -2455,25 +2455,28 @@
 re();%0A      
-T
+  _t
 argetEncoder
@@ -2481,36 +2481,17 @@
 erModel.
-TargetEncoderOutput 
+_
 output =
@@ -2603,119 +2603,190 @@
     
-_targetEncoderModel = new TargetEncoderModel(_result, _parms, output, tec);%0A      DKV.put(_targetEncoderModel);
+  _job.update(1);%0A        _targetEncoderModel.update(_job);%0A      %7D finally %7B%0A        if (_targetEncoderModel != null) %7B%0A          _targetEncoderModel.unlock(_job);%0A        %7D%0A      %7D
 %0A   

Fix grid search for Target Encoder