package edu.mit.csail.db.ml.server.algorithm.similarity;

import edu.mit.csail.db.ml.server.algorithm.similarity.comparators.*;
import edu.mit.csail.db.ml.server.storage.TransformerDao;
import jooq.sqlite.gen.Tables;
import jooq.sqlite.gen.tables.records.LinearmodelRecord;
import modeldb.BadRequestException;
import modeldb.ModelCompMetric;
import modeldb.ResourceNotFoundException;
import org.jooq.DSLContext;
import org.jooq.TableField;

import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

public class SimilarModels {
  // TODO: Should this go into the conf file?
  private static final int MODEL_LIMIT = 1000;

  private static ModelComparator comparatorForCompMetric(ModelCompMetric compMetric) {
    switch (compMetric) {
      case MODEL_TYPE: return new TransformerTypeComparator();
      case PROBLEM_TYPE: return new ProblemTypeComparator();
      case EXPERIMENT_RUN: return new ExperimentRunComparator();
      case PROJECT: return new ProjectComparator();
      case RMSE: return new LinearModelComparator() {
        @Override
        public TableField<LinearmodelRecord, Double> getComparisonField() {
          return Tables.LINEARMODEL.RMSE;
        }
      };
      case EXPLAINED_VARIANCE: return new LinearModelComparator() {
        @Override
        public TableField<LinearmodelRecord, Double> getComparisonField() {
          return Tables.LINEARMODEL.EXPLAINEDVARIANCE;
        }
      };
      case R2: return new LinearModelComparator() {
        @Override
        public TableField<LinearmodelRecord, Double> getComparisonField() {
          return Tables.LINEARMODEL.R2;
        }
      };
      default: return new DummyComparator();
    }
  }

  public static List<Integer> similarModels(
    int modelId,
    List<ModelCompMetric> compMetrics,
    int numModels,
    DSLContext ctx
  ) throws ResourceNotFoundException, BadRequestException {
    if (!TransformerDao.exists(modelId, ctx)) {
      throw new ResourceNotFoundException(String.format(
        "Cannot find models similar to Transformer %d because that Transformer does not exist",
        modelId
      ));
    }
    if (numModels <= 0) {
      throw new BadRequestException(String.format(
        "You can't ask for %d similar models, you must request a positive number of models",
        numModels
      ));
    }
    List<Integer> resultingModels = compMetrics
      .stream()
      .map(SimilarModels::comparatorForCompMetric)
      .reduce(
        Collections.<Integer>emptyList(),
        (partial, comparator) -> comparator.similarModels(modelId, partial, MODEL_LIMIT, ctx),
        (oldModels, newModels) -> newModels
      )
      .stream()
      .filter(s -> !s.equals(modelId))
      .collect(Collectors.toList());
    return resultingModels.subList(0, Math.min(resultingModels.size(), numModels));
  }
}

@@ -524,16 +524,128 @@
 ctors;%0A%0A
+/**%0A * This exposes the similarModels method that allows finding models similar to a model with a given ID.%0A */%0A
 public c
@@ -759,16 +759,246 @@
  1000;%0A%0A
+  /**%0A   * This creates the appropriate ModelComparator object for a given comparison metric.%0A   * @param compMetric - The model comparison metric.%0A   * @return The ModelComparator object that corresponds to the compMetric.%0A   */%0A
   privat
@@ -2049,16 +2049,778 @@
  %7D%0A  %7D%0A%0A
+  /**%0A   * Find the models similar to the model with the given ID.%0A   * @param modelId - The ID of the given model. This method will find models similar to this model.%0A   * @param compMetrics - The metrics by which to measure similarity. The first metric is the most important and thus%0A   *                    models that are most similar according to this metric will be first in the returned list. Each%0A   *                    successive metric is used to break ties in the previous metric.%0A   * @param numModels - The maximum number of models that should be returned.%0A   * @param ctx - The context for accessing the database.%0A   * @return The IDs of similar models, with the most similar model first and successively less similar%0A   * models following.%0A   */%0A
   public
@@ -3010,16 +3010,54 @@
 ption %7B%0A
+    // Fail if modelID doesn't exist.%0A
     if (
@@ -3276,24 +3276,69 @@
    ));%0A    %7D
+%0A%0A    // Fail if number of models is invalid.
 %0A    if (num
@@ -3525,24 +3525,57 @@
    ));%0A    %7D
+%0A%0A    // Find the similar models.
 %0A    List%3CIn
@@ -3937,16 +3937,16 @@
 delId))%0A
-
       .c
@@ -3973,16 +3973,57 @@
 List());
+%0A%0A    // Return at most numModels models.
 %0A    ret

Comment similarModels