package com.airbnb.aerosolve.core.function;

import com.airbnb.aerosolve.core.ModelRecord;
import com.airbnb.aerosolve.core.NDTreeNode;
import com.airbnb.aerosolve.core.models.NDTreeModel;
import lombok.extern.slf4j.Slf4j;

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

@Slf4j
public class MultiDimensionSpline implements Function {
  private static final long serialVersionUID = 5166347177557769302L;
  private NDTreeModel ndTreeModel;
  // NDTree leaf maps to spline point
  private Map<Integer, List<MultiDimensionPoint>> weights;

  public MultiDimensionSpline(NDTreeModel ndTreeModel) {
    this.ndTreeModel = ndTreeModel;
    Map<List<Float>, MultiDimensionPoint> points = new HashMap<>();
    weights = new HashMap<>();

    NDTreeNode[] nodes = ndTreeModel.getNodes();
    for (int i = 0; i < nodes.length; i++) {
      NDTreeNode node = nodes[i];
      if (node.getCoordinateIndex() == NDTreeModel.LEAF) {
        List<MultiDimensionPoint> list = MultiDimensionPoint.getCombinationWithoutDuplication(
            node.getMin(), node.getMax(), points);
        if (list != null && !list.isEmpty()) {
          weights.put(i, list);
        } else {
          log.info("leaf node return no MultiDimensionPoint {}", node);
        }
      }
    }
  }

  @Override
  public Function aggregate(Iterable<Function> functions, float scale, int numBins) {
    return null;
  }

  @Override
  public float evaluate(float ... coordinates) {
    List<MultiDimensionPoint> list = getNearbyPoints(coordinates);
    double[] distance = new double[list.size()];
    double sum = 0;
    for (int i = 0; i < list.size(); i++) {
      MultiDimensionPoint point = list.get(i);
      distance[i] = point.getDistance(coordinates);
      sum += distance[i];
    }
    float score = 0;
    for (int i = 0; i < list.size(); i++) {
      MultiDimensionPoint point = list.get(i);
      score += point.getWeight() * (distance[i]/sum);
    }
    return score;
  }

  @Override
  public void update(float delta, float ... values) {
    List<MultiDimensionPoint> list = getNearbyPoints(values);
    double[] distance = new double[list.size()];
    double sum = 0;
    for (int i = 0; i < list.size(); i++) {
      MultiDimensionPoint point = list.get(i);
      distance[i] = point.getDistance(values);
      sum += distance[i];
    }
    for (int i = 0; i < list.size(); i++) {
      MultiDimensionPoint point = list.get(i);
      point.updateWeight(delta * (distance[i]/sum));
    }
  }

  @Override
  public ModelRecord toModelRecord(String featureFamily, String featureName) {
    return null;
  }

  @Override
  public void setPriors(float[] params) {

  }

  @Override
  public void LInfinityCap(float cap) {

  }

  @Override
  public float LInfinityNorm() {
    return 0;
  }

  private List<MultiDimensionPoint> getNearbyPoints(float ... coordinates) {
    int index = ndTreeModel.leaf(coordinates);
    assert (index != -1 && weights.containsKey(index));
    return weights.get(index);
  }
}

@@ -238,60 +238,9 @@
 til.
-HashMap;%0Aimport java.util.List;%0Aimport java.util.Map
+*
 ;%0A%0A@
@@ -380,16 +380,22 @@
 private 
+final 
 NDTreeMo
@@ -459,16 +459,22 @@
 private 
+final 
 Map%3CInte
@@ -513,16 +513,66 @@
 weights;
+%0A  private final List%3CMultiDimensionPoint%3E points;
 %0A%0A  publ
@@ -706,16 +706,19 @@
 %3E points
+Map
  = new H
@@ -1091,16 +1091,19 @@
 , points
+Map
 );%0A     
@@ -1295,124 +1295,641 @@
 %7D%0A  
-%7D%0A%0A  @Override%0A  public Function aggregate(Iterable%3CFunction%3E functions, float scale, int numBins) %7B%0A    return null
+  points = new ArrayList%3C%3E(pointsMap.values());%0A  %7D%0A%0A  @Override // it doesn't need numBins just like linear%0A  public Function aggregate(Iterable%3CFunction%3E functions, float scale, int numBins) %7B%0A    // functions size == 1/scale%0A    int length = points.size();%0A    float%5B%5D aggWeights = new float%5Blength%5D;%0A    for (Function fun: functions) %7B%0A      MultiDimensionSpline spline = (MultiDimensionSpline) fun;%0A%0A      for (int i = 0; i %3C length; i++) %7B%0A        aggWeights%5Bi%5D += scale * spline.points.get(i).getWeight();%0A      %7D%0A    %7D%0A    for (int i = 0; i %3C length; i++) %7B%0A      points.get(i).setWeight(aggWeights%5Bi%5D);%0A    %7D%0A    return this
 ;%0A  

add aggregate for multi dimension spline