package hex.genmodel.algos.glm;

import com.sun.corba.se.spi.orb.StringPair;
import hex.genmodel.MojoModel;

import java.util.HashMap;
import java.util.stream.IntStream;

abstract class GlmMojoModelBase extends MojoModel {

  boolean _useAllFactorLevels;
  public StringPair[] _interaction_pairs;
  HashMap<String, Integer[]> _interaction_mapping;

  int _cats;
  int[] _catModes;
  int[] _catOffsets;

  int _nums;
  double[] _numMeans;
  boolean _meanImputation;

  double[] _beta;

  String _family;

  GlmMojoModelBase(String[] columns, String[][] domains, String responseColumn) {
    super(columns, domains, responseColumn);
    mapInteractions();
  }

  void init() { /* do nothing by default */ }

  @Override
  public final double[] score0(double[] data, double[] preds) {
    if (_meanImputation)
      imputeMissingWithMeans(data);

    return glmScore0(data, preds);
  }

  abstract double[] glmScore0(double[] data, double[] preds);

  private void imputeMissingWithMeans(double[] data) {
    for (int i = 0; i < _cats; ++i) {
      if (Double.isNaN(data[i])) {
        if (isInteraction(_names[i])) {
          imputeInteraction(data, i);
        } else {
          data[i] = _catModes[i];
        }
      }
    }
    for (int i = 0; i < _nums; ++i) {
      if (Double.isNaN(data[i + _cats])) {
        if (isInteraction(_names[i + _cats])) {
          imputeInteraction(data, i + _cats);
        } else {
          data[i + _cats] = _numMeans[i];  
        }
      } 
    }
  }

  private void mapInteractions() {
    if (_interaction_pairs == null) {
      _interaction_mapping = null;
      return;
    }
    
    _interaction_mapping = new HashMap<String, Integer[]>();
    for (StringPair pair: _interaction_pairs) {
      int index1 = IntStream.range(0, _names.length)
              .filter(i -> pair.getFirst() == _names[i])
              .findFirst()
              .orElse(-1);
      int index2 = IntStream.range(0, _names.length)
              .filter(i -> pair.getSecond() == _names[i])
              .findFirst()
              .orElse(-1);
      _interaction_mapping.put((pair.getFirst() + "_" + pair.getSecond()), new Integer[]{index1, index2});
    }
  }
  
  private boolean isInteraction(String name) {
    if(_interaction_pairs == null) {return false;}
    return _interaction_mapping.containsKey(name);
  }
  
  private void imputeInteraction(double[] data, int index) {
    Integer[] pairs = _interaction_mapping.get(_names[index]);
    if (_domains[pairs[0]] != null && _domains[pairs[1]] != null) {
      data[index] = (_domains[pairs[0]].length * data[pairs[0]]) + data[pairs[1]] - 1;
    } else if (_domains[pairs[0]] == null && _domains[pairs[1]] == null) {
      data[index] = data[pairs[0]] * data[pairs[1]];
    } else {
      data[index] = _domains[pairs[0]] == null ? data[pairs[0]] : data[pairs[1]];  
    }
  }
  
  @Override
  public String[] getOutputNames() {
    // special handling of binomial case where response domain is not represented
    if (nclasses() == 2 && getDomainValues(getResponseIdx()) == null) {
      return new String[]{"predict", "0", "1"};
    }
    return super.getOutputNames();
  }
}

@@ -106,70 +106,8 @@
 l;%0A%0A
-import java.util.HashMap;%0Aimport java.util.stream.IntStream;%0A%0A
 abst
@@ -231,59 +231,8 @@
 irs;
-%0A  HashMap%3CString, Integer%5B%5D%3E _interaction_mapping;
 %0A%0A  
@@ -515,31 +515,8 @@
 n);%0A
-    mapInteractions();%0A
   %7D%0A
@@ -890,26 +890,24 @@
  _cats; ++i)
- %7B
 %0A      if (D
@@ -931,115 +931,8 @@
 i%5D))
- %7B%0A        if (isInteraction(_names%5Bi%5D)) %7B%0A          imputeInteraction(data, i);%0A        %7D else %7B%0A         
  dat
@@ -956,32 +956,8 @@
 i%5D;%0A
-        %7D%0A      %7D%0A    %7D%0A
     
@@ -987,18 +987,16 @@
 ms; ++i)
- %7B
 %0A      i
@@ -1033,1558 +1033,44 @@
 %5D)) 
-%7B%0A        if (isInteraction(_names%5Bi + _cats%5D)) %7B%0A          imputeInteraction(data, i + _cats);%0A        %7D else %7B%0A          data%5Bi + _cats%5D = _numMeans%5Bi%5D;  %0A        %7D%0A      %7D %0A    %7D%0A  %7D%0A%0A  private void mapInteractions() %7B%0A    if (_interaction_pairs == null) %7B%0A      _interaction_mapping = null;%0A      return;%0A    %7D%0A    %0A    _interaction_mapping = new HashMap%3CString, Integer%5B%5D%3E();%0A    for (StringPair pair: _interaction_pairs) %7B%0A      int index1 = IntStream.range(0, _names.length)%0A              .filter(i -%3E pair.getFirst() == _names%5Bi%5D)%0A              .findFirst()%0A              .orElse(-1);%0A      int index2 = IntStream.range(0, _names.length)%0A              .filter(i -%3E pair.getSecond() == _names%5Bi%5D)%0A              .findFirst()%0A              .orElse(-1);%0A      _interaction_mapping.put((pair.getFirst() + %22_%22 + pair.getSecond()), new Integer%5B%5D%7Bindex1, index2%7D);%0A    %7D%0A  %7D%0A  %0A  private boolean isInteraction(String name) %7B%0A    if(_interaction_pairs == null) %7Breturn false;%7D%0A    return _interaction_mapping.containsKey(name);%0A  %7D%0A  %0A  private void imputeInteraction(double%5B%5D data, int index) %7B%0A    Integer%5B%5D pairs = _interaction_mapping.get(_names%5Bindex%5D);%0A    if (_domains%5Bpairs%5B0%5D%5D != null && _domains%5Bpairs%5B1%5D%5D != null) %7B%0A      data%5Bindex%5D = (_domains%5Bpairs%5B0%5D%5D.length * data%5Bpairs%5B0%5D%5D) + data%5Bpairs%5B1%5D%5D - 1;%0A    %7D else if (_domains%5Bpairs%5B0%5D%5D == null && _domains%5Bpairs%5B1%5D%5D == null) %7B%0A      data%5Bindex%5D = data%5Bpairs%5B0%5D%5D * data%5Bpairs%5B1%5D%5D;%0A    %7D else %7B%0A      data%5Bindex%5D = _domains%5Bpairs%5B0%5D%5D == null ? data%5Bpairs%5B0%5D%5D : data%5Bpairs%5B1%5D%5D;  %0A    %7D
+data%5Bi + _cats%5D = _numMeans%5Bi%5D;
 %0A  %7D%0A
-  
 %0A  @

Revert "Add impute interactions to GLM mojo model base"