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

import jooq.sqlite.gen.Tables;
import jooq.sqlite.gen.tables.records.DataframeRecord;
import jooq.sqlite.gen.tables.records.DataframesplitRecord;
import jooq.sqlite.gen.tables.records.EventRecord;
import jooq.sqlite.gen.tables.records.RandomspliteventRecord;
import modeldb.RandomSplitEvent;
import modeldb.RandomSplitEventResponse;
import org.jooq.DSLContext;

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

public class RandomSplitEventDao {
  public static RandomSplitEventResponse store(RandomSplitEvent rse, DSLContext ctx) {
    DataframeRecord oldDf = DataFrameDao.store(rse.oldDataFrame, rse.experimentRunId, ctx);


    RandomspliteventRecord rseRec = ctx.newRecord(Tables.RANDOMSPLITEVENT);
    rseRec.setId(null);
    rseRec.setInputdataframeid(oldDf.getId());
    rseRec.setRandomseed(rse.seed);
    rseRec.setExperimentrun(rse.experimentRunId);
    rseRec.store();

    EventRecord ev = EventDao.store(rseRec.getId(), "random split", rse.experimentRunId, ctx);

    List<DataframeRecord> splitDfs = rse
      .splitDataFrames
      .stream()
      .map(df -> DataFrameDao.store(df, rse.experimentRunId, ctx))
      .collect(Collectors.toList());

    List<Integer> splitIds = IntStream
      .range(0, rse.splitDataFrames.size())
      .map(ind -> {
        DataframesplitRecord splRec = ctx.newRecord(Tables.DATAFRAMESPLIT);
        splRec.setId(null);
        splRec.setSpliteventid(rseRec.getId());
        splRec.setWeight(rse.weights.get(ind).floatValue());
        splRec.setDataframeid(splitDfs.get(ind).getId());
        splRec.setExperimentrun(rse.experimentRunId);
        splRec.store();
        return splRec.getId();
      })
      .boxed()
      .collect(Collectors.toList());

    return new RandomSplitEventResponse(oldDf.getId(), splitIds, ev.getId());
  }
}

@@ -316,95 +316,69 @@
 ldb.
-RandomSplitEvent;%0Aimport modeldb.RandomSplitEventResponse;%0Aimport org.jooq.DSLContext;%0A
+*;%0Aimport org.jooq.DSLContext;%0A%0Aimport java.util.Collections;
 %0Aimp
@@ -1761,24 +1761,782 @@
 toList());%0A%0A
+    // Store a TransformEvent for each split. This allows us to preserve the ancestor chain of DataFrames.%0A    DataFrame oldDataFrame = rse.oldDataFrame.setId(oldDf.getId());%0A    Transformer rseTransformer = new Transformer(-1, Collections.emptyList(), %22RandomSplitTransformer%22, %22%22);%0A    IntStream.range(0, rse.splitDataFrames.size())%0A      .forEach(index -%3E %7B%0A        DataFrame splitDataFrame = rse.splitDataFrames.get(index).setId(splitDfs.get(index).getId());%0A        TransformEvent te = new TransformEvent(%0A          oldDataFrame,%0A          splitDataFrame,%0A          rseTransformer,%0A          Collections.emptyList(),%0A          Collections.emptyList(),%0A          rse.experimentRunId%0A        );%0A        TransformEventDao.store(te, ctx, false);%0A      %7D);%0A%0A
     return n

Store TransformEvents for each split in a RandomSplitEvent