package edu.mit.csail.db.ml.server.storage;

import jooq.sqlite.gen.Tables;
import jooq.sqlite.gen.tables.records.*;
import modeldb.*;
import org.jooq.DSLContext;

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

public class CrossValidationEventDao {
  public static CrossValidationEventResponse store(CrossValidationEvent cve, DSLContext ctx) {
    DataframeRecord df = DataFrameDao.store(cve.df, cve.experimentRunId, ctx);
    TransformerspecRecord spec = TransformerSpecDao.store(cve.spec, cve.experimentRunId, ctx);

    CrossvalidationeventRecord cveRec = ctx.newRecord(Tables.CROSSVALIDATIONEVENT);
    cveRec.setId(null);
    cveRec.setDf(df.getId());
    cveRec.setSpec(spec.getId());
    cveRec.setNumfolds(cve.folds.size());
    cveRec.setRandomseed(cve.seed);
    cveRec.setEvaluator(cve.evaluator);
    cveRec.setExperimentrun(cve.experimentRunId);
    cveRec.store();

    EventRecord ev = EventDao.store(cveRec.getId(), "cross validation", cve.experimentRunId, ctx);

    List<FitEventResponse> fers = cve
      .folds
      .stream()
      .map(fold -> new FitEvent(
        fold.trainingDf,
        cve.spec.setId(spec.getId()),
        fold.model,
        cve.featureColumns,
        cve.predictionColumns,
        cve.labelColumns,
        cve.experimentRunId
      ).setProblemType(cve.problemType))
      .map(fe -> FitEventDao.store(fe, ctx))
      .collect(Collectors.toList());

    List<MetricEventResponse> mers = IntStream
      .range(0, cve.folds.size())
      .mapToObj(ind -> new MetricEvent(
        cve.folds.get(ind).validationDf,
        cve.folds.get(ind).model.setId(fers.get(ind).modelId),
        cve.evaluator,
        cve.folds.get(ind).score,
        cve.labelColumns.get(0), // TODO: Are these reasonable settings?
        cve.predictionColumns.get(0),
        cve.experimentRunId
      ))
      .map(me -> MetricEventDao.store(me, ctx))
      .collect(Collectors.toList());


    List<CrossValidationFoldResponse> foldResponses = IntStream
      .range(0, cve.folds.size())
      .mapToObj(ind -> new CrossValidationFoldResponse(
        fers.get(ind).modelId,
        mers.get(ind).dfId,
        fers.get(ind).dfId
      ))
      .collect(Collectors.toList());

    mers.forEach(me -> {
      CrossvalidationfoldRecord rec = ctx.newRecord(Tables.CROSSVALIDATIONFOLD);
      rec.setId(null);
      rec.setMetric(me.getMetricEventId());
      rec.setEvent(cveRec.getId());
      rec.setExperimentrun(cve.experimentRunId);
      rec.store();
      rec.getId();
    });

    return new CrossValidationEventResponse(df.getId(), spec.getId(), ev.getId(), foldResponses, cveRec.getId());
  }
}

@@ -256,46 +256,351 @@
 m;%0A%0A
-public class CrossValidationEventDao %7B
+/**%0A * This class contains logic for storing and reading CrossValidationEvents.%0A */%0Apublic class CrossValidationEventDao %7B%0A  /**%0A   * Store the given CrossValidationEvent in the database.%0A   * @param cve - The CrossValidationEvent.%0A   * @param ctx - The database context.%0A   * @return The response after storing the CrossValidationEvent.%0A   */
 %0A  p
@@ -691,16 +691,91 @@
  ctx) %7B%0A
+    // Store the training DataFrame that is used for CrossValidationEvent.%0A
     Data
@@ -844,16 +844,96 @@
 d, ctx);
+%0A%0A    // Store the TransformerSpec that is being evaluated for cross validation.
 %0A    Tra
@@ -1017,24 +1017,81 @@
 nId, ctx);%0A%0A
+    // Store an entry in the CrossValidationEvent table.%0A
     Crossval
@@ -1435,24 +1435,66 @@
 c.store();%0A%0A
+    // Store an entry in the Event table.%0A
     EventRec
@@ -1577,24 +1577,63 @@
 nId, ctx);%0A%0A
+    // Store a FitEvent for each fold.%0A
     List%3CFit
@@ -2036,24 +2036,66 @@
 toList());%0A%0A
+    // Store a MetricEvent for each fold.%0A
     List%3CMet
@@ -2598,16 +2598,87 @@
 st());%0A%0A
+    // Prepare the responses that will be returned (one for each fold).
 %0A    Lis
@@ -2953,24 +2953,74 @@
 toList());%0A%0A
+    // Store a CrossValidationFold for each fold.%0A
     mers.for
@@ -3308,24 +3308,90 @@
 );%0A    %7D);%0A%0A
+    // Return the response containing the IDs of the stored rows.%0A
     return n

Store CrossValidationEventDao