package edu.umass.cs.jfoley.coop.index.covariate;

import ciir.jfoley.chai.collections.Pair;
import ciir.jfoley.chai.io.Directory;
import edu.umass.cs.ciir.waltz.coders.Coder;
import edu.umass.cs.ciir.waltz.coders.data.BufferList;
import edu.umass.cs.ciir.waltz.coders.data.DataChunk;
import edu.umass.cs.ciir.waltz.coders.kinds.DeltaIntListCoder;
import edu.umass.cs.ciir.waltz.galago.io.GalagoIO;
import edu.umass.cs.jfoley.coop.document.CoopDoc;
import edu.umass.cs.jfoley.coop.document.DocVar;
import edu.umass.cs.jfoley.coop.index.general.DocumentSetWriter;
import edu.umass.cs.jfoley.coop.index.general.IndexItemWriter;
import edu.umass.cs.jfoley.coop.schema.DocVarSchema;
import edu.umass.cs.jfoley.coop.schema.IndexConfiguration;
import org.lemurproject.galago.utility.Parameters;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.io.InputStream;

/**
 * Given two variables {@link DocVar}, create an IOMap that will allow for somewhat efficient cell-based lookup of documents that belong in that cell.
 * @author jfoley
 */
public class MapCovariateSpaceWriter<A extends Comparable<A>,B extends Comparable<B>> extends IndexItemWriter {

  private final DocVarSchema<A> xSchema;
  private final DocVarSchema<B> ySchema;
  private final Coder<A> xCoder;
  private final Coder<B> yCoder;
  private final DocumentSetWriter<Covariable> writer;

  public final class Covariable extends Pair<A,B> implements Comparable<Covariable> {
    public Covariable(A left, B right) {
      super(left, right);
    }

    @Override
    public int compareTo(@Nonnull Covariable o) {
      int cmp = left.compareTo(o.left);
      if(cmp != 0) return cmp;

      return right.compareTo(o.right);
    }
  }
  public final class CovariableCoder extends Coder<Covariable> {
    @Override
    public boolean knowsOwnSize() {
      return true;
    }

    @Nonnull
    @Override
    public DataChunk writeImpl(Covariable obj) throws IOException {
      BufferList bl = new BufferList();
      bl.add(xCoder, obj.left);
      bl.add(yCoder, obj.right);
      return bl.compact();
    }

    @Nonnull
    @Override
    public Covariable readImpl(InputStream inputStream) throws IOException {
      A left = xCoder.readImpl(inputStream);
      B right = yCoder.readImpl(inputStream);
      return new Covariable(left, right);
    }
  }

  public MapCovariateSpaceWriter(Directory outputDir, IndexConfiguration cfg, DocVarSchema<A> xSchema, DocVarSchema<B> ySchema) throws IOException {
    super(outputDir, cfg);
    this.xSchema = xSchema;
    this.ySchema = ySchema;
    this.writer = new DocumentSetWriter<>(
        GalagoIO.getIOMapWriter(
            new CovariableCoder(),
            new DeltaIntListCoder(),
            outputDir.childPath("covar." + xSchema.getName() + "." + ySchema.getName()), // TODO, do variables need short-names?
            Parameters.parseArray(
                "covarXSchema", xSchema.toJSON(),
                "covarYSchema", ySchema.toJSON()
            )
        )
    );
    this.xCoder = xSchema.getCoder().lengthSafe();
    this.yCoder = ySchema.getCoder().lengthSafe();
  }

  @Override
  public void process(CoopDoc document) throws IOException {
    // Skip documents that are missing either covariate.
    A varA = document.getVariableValue(xSchema);
    if(varA == null) return;
    B varB = document.getVariableValue(ySchema);
    if(varB == null) return;

    Covariable covar = new Covariable(varA, varB);
    writer.process(covar, document.getIdentifier());
  }

  @Override
  public void close() throws IOException {
    writer.close();
  }
}

@@ -1249,74 +1249,8 @@
 ma;%0A
-  private final Coder%3CA%3E xCoder;%0A  private final Coder%3CB%3E yCoder;%0A
   pr
@@ -1289,16 +1289,21 @@
 variable
+%3CA,B%3E
 %3E writer
@@ -1306,37 +1306,38 @@
 iter;%0A%0A  public 
-final
+static
  class Covariabl
@@ -1337,16 +1337,66 @@
 variable
+%3CA extends Comparable%3CA%3E, B extends Comparable%3CB%3E%3E
  extends
@@ -1430,32 +1430,37 @@
 rable%3CCovariable
+%3CA,B%3E
 %3E %7B%0A    public C
@@ -1579,16 +1579,21 @@
 variable
+%3CA,B%3E
  o) %7B%0A  
@@ -1720,21 +1720,22 @@
  public 
-final
+static
  class C
@@ -1752,36 +1752,276 @@
 oder
- extends Coder%3CCovariable%3E %7B
+%3CA extends Comparable%3CA%3E, B extends Comparable%3CB%3E%3E extends Coder%3CCovariable%3CA,B%3E%3E %7B%0A    private Coder%3CA%3E xCoder;%0A    private Coder%3CB%3E yCoder;%0A%0A    public CovariableCoder(Coder%3CA%3E xCoder, Coder%3CB%3E yCoder) %7B%0A      this.xCoder = xCoder;%0A      this.yCoder = yCoder;%0A    %7D%0A
 %0A   
@@ -2161,16 +2161,21 @@
 variable
+%3CA,B%3E
  obj) th
@@ -2380,16 +2380,21 @@
 variable
+%3CA,B%3E
  readImp
@@ -2559,16 +2559,18 @@
 variable
+%3C%3E
 (left, r
@@ -2811,32 +2811,142 @@
 hema = ySchema;%0A
+    Coder%3CA%3E xCoder = xSchema.getCoder().lengthSafe();%0A    Coder%3CB%3E yCoder = ySchema.getCoder().lengthSafe();%0A
     this.writer 
@@ -3032,25 +3032,41 @@
 ariableCoder
-(
+%3C%3E(xCoder, yCoder
 ),%0A         
@@ -3391,110 +3391,8 @@
  );%0A
-    this.xCoder = xSchema.getCoder().lengthSafe();%0A    this.yCoder = ySchema.getCoder().lengthSafe();%0A
   %7D%0A
@@ -3693,16 +3693,21 @@
 variable
+%3CA,B%3E
  covar =
@@ -3721,16 +3721,18 @@
 variable
+%3C%3E
 (varA, v

Make static and make generic explicit