package water.rapids.ast.prims.time;

import org.joda.time.DateTimeZone;
import org.joda.time.IllegalFieldValueException;
import org.joda.time.MutableDateTime;
import water.Key;
import water.MRTask;
import water.fvec.Chunk;
import water.fvec.Frame;
import water.fvec.NewChunk;
import water.fvec.Vec;
import water.rapids.Val;
import water.rapids.ast.AstBuiltin;
import water.rapids.vals.ValFrame;
import water.util.ArrayUtils;

import java.util.ArrayList;


/**
 * Convert year, month, day, hour, minute, sec, msec to Unix epoch time
 * (in milliseconds).
 *
 * This is a replacement for {@code AstMktime} class.
 */
public class AstMoment extends AstBuiltin<AstMoment> {

  @Override public int nargs() {
    return 8;
  }

  public String[] args() {
    return new String[]{"yr", "mo", "dy", "hr", "mi", "se", "ms"};
  }

  @Override public String str() {
    return "moment";
  }


  @Override
  protected ValFrame exec(Val[] args) {

    // Parse the input arguments, verifying their validity.
    boolean naResult = false;
    long numRows = -1;
    int[] timeparts = new int[7];
    ArrayList<Integer> chunksmap = new ArrayList<>(7);
    ArrayList<Vec> timevecs = new ArrayList<>(7);
    for (int i = 0; i < 7; i++) {
      Val vi = args[i + 1];
      if (vi.isFrame()) {
        Frame fr = vi.getFrame();
        if (fr.numCols() != 1)
          throw new IllegalArgumentException("Argument " + i + " is a frame with " + fr.numCols() + " columns");
        if (!fr.vec(0).isNumeric())
          throw new IllegalArgumentException("Argument " + i + " is not a numeric column");
        if (fr.numRows() == 0)
          throw new IllegalArgumentException("Column " + i + " has 0 rows");
        if (fr.numRows() == 1) {
          double d = fr.vec(0).at(0);
          if (Double.isNaN(d))
            naResult = true;
          else
            timeparts[i] = (int) d;
        } else {
          if (numRows == -1)
            numRows = fr.numRows();
          if (fr.numRows() != numRows)
            throw new IllegalArgumentException("Incompatible vec " + i + " having " + fr.numRows() + " rows, whereas " +
                                               "other vecs have " + numRows + " rows.");
          timevecs.add(fr.vec(0));
          chunksmap.add(i);
        }
      } else if (vi.isNum()){
        double d = vi.getNum();
        if (Double.isNaN(d))
          naResult = true;
        else
          timeparts[i] = (int) d;
      } else {
        throw new IllegalArgumentException("Argument " + i + " is neither a number nor a frame");
      }
    }

    // If all arguments are scalars, return a 1x1 frame
    if (timevecs.isEmpty()) {
      double val = Double.NaN;
      if (!naResult) {
        try {
          val = new MutableDateTime(timeparts[0], timeparts[1], timeparts[2], timeparts[3], timeparts[4],
                                    timeparts[5], timeparts[6], DateTimeZone.UTC).getMillis();
        } catch (IllegalFieldValueException ignored) {}
      }
      return make1x1Frame(val);
    }

    // Some arguments are vecs -- create a frame of the same size
    Vec[] vecs = timevecs.toArray(new Vec[timevecs.size()]);
    int[] cm = ArrayUtils.toPrimitive(chunksmap);
    Frame fr = new SetTimeTask(timeparts, cm, naResult)
        .doAll(Vec.T_TIME, vecs)
        .outputFrame(Key.<Frame>make(), new String[]{"time"}, null);

    return new ValFrame(fr);
  }

  private ValFrame make1x1Frame(double val) {
    Vec v = Vec.makeTimeVec(new double[]{val}, null);
    Frame f = new Frame(new String[]{"time"}, new Vec[]{v});
    return new ValFrame(f);
  }


  private static class SetTimeTask extends MRTask<SetTimeTask> {
    private int[] tp;
    private int[] cm;
    private boolean na;

    /**
     * @param timeparts is the array of [year, month, day, hrs, mins, secs, ms]
     *                  for all constant parts of the date;
     * @param chunksmap is a mapping between chunks indices and the timeparts
     *                  array. For example, if {@code chunksmap = [1, 2]},
     *                  then the first chunk describes the "month" part of the
     *                  date, and the second chunk the "day" part.
     */
    public SetTimeTask(int[] timeparts, int[] chunksmap, boolean naResult) {
      tp = timeparts;
      cm = chunksmap;
      na = naResult;
    }

    @Override public void map(Chunk[] chks, NewChunk nc) {
      int nVecs = cm.length;
      assert chks.length == nVecs;
      MutableDateTime dt = new MutableDateTime(0, DateTimeZone.UTC);
      int nChunkRows = chks[0]._len;
      if (na) {
        for (int i = 0; i < nChunkRows; i++) {
          nc.addNum(Double.NaN);
        }
      } else {
        BYROW:
        for (int i = 0; i < nChunkRows; i++) {
          for (int j = 0; j < nVecs; j++) {
            double d = chks[j].atd(i);
            if (Double.isNaN(d)) {
              nc.addNum(Double.NaN);
              continue BYROW;
            }
            tp[cm[j]] = (int) d;
          }
          try {
            dt.setDateTime(tp[0], tp[1], tp[2], tp[3], tp[4], tp[5], tp[6]);
            nc.addNum(dt.getMillis());
          } catch (IllegalFieldValueException e) {
            nc.addNum(Double.NaN);
          }
        }
      }
    }
  }
}

@@ -52,28 +52,26 @@
 da.time.
-DateTimeZone
+Chronology
 ;%0Aimport
@@ -134,31 +134,36 @@
 da.time.
-MutableDateTime
+chrono.ISOChronology
 ;%0Aimport
@@ -2736,35 +2736,58 @@
   val = 
-new Mutable
+ISOChronology.getInstanceUTC().get
 DateTime
 (timepar
@@ -2778,16 +2778,22 @@
 DateTime
+Millis
 (timepar
@@ -2826,16 +2826,30 @@
 arts%5B2%5D,
+%0A             
  timepar
@@ -2872,101 +2872,35 @@
 %5B4%5D,
-%0A                                    timeparts%5B5%5D, timeparts%5B6%5D, DateTimeZone.UTC).getMillis(
+ timeparts%5B5%5D, timeparts%5B6%5D
 );%0A 
@@ -4437,68 +4437,61 @@
     
-MutableDateTime dt = new MutableDateTime(0, DateTimeZone.
+Chronology chronology = ISOChronology.getInstance
 UTC
+(
 );%0A 
@@ -4986,11 +4986,35 @@
    d
-t.s
+ouble millis = chronology.g
 etDa
@@ -5019,16 +5019,22 @@
 DateTime
+Millis
 (tp%5B0%5D, 
@@ -5102,22 +5102,14 @@
 Num(
-dt.getM
+m
 illis
-()
 );%0A 

Replace use of MutableDateTime with Chronology (avoid 1 level of indirection)