#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
#-------------------------------------------------------------
# source("scripts/pipelines/scripts/bandit.dml") as hb

F = read($1, data_type="frame", format="csv", header=FALSE, 
  naStrings= ["NA", "null","  ","NaN", "nan", "", "?", "99999"]);
metaInfo = read($2, data_type="frame", format="csv", header=FALSE);
primitives = read($3, data_type = "frame", format="csv", header= TRUE)
param = read($4, data_type = "frame", format="csv", header= TRUE)
sample = $5
topK = $6
resources = $7
crossValidations = $8
weightedAccuracy = $9
logical = read($11, data_type = "frame", format="csv", header= FALSE)

[result, res] = startCleaning(F=F, logical=logical, target="classification",  metaInfo=metaInfo, 
  primitives=primitives, param=param, k=topK, sample=sample, isWeighted=weightedAccuracy, 
  R=resources, cv=crossValidations, verbose=TRUE)
  
output = as.logical(as.scalar(res[1,1] < res[1,2]))
write(output, $10, format = "text")





startCleaning = function(Frame[Unknown] F, Frame[Unknown] logical, String target = "classification", 
  Frame[Unknown] metaInfo, Frame[Unknown] primitives, Frame[Unknown] param, Integer k, Double sample,
  Boolean isWeighted = TRUE, Integer R=50, Integer cv=3, Boolean verbose = FALSE)
  return (Frame[Unknown] result, Matrix[Double] res)
{
  
  res = as.matrix(0)
  result = as.frame(0)
  pip = as.frame("")
  hp = as.matrix(0)
  acc = as.matrix(0)
  
  if(nrow(metaInfo) < 2)
    stop("incomplete meta info")
    
  # initialize variables
  eX = as.matrix(0)
  eY = as.matrix(0)
  
  getSchema = metaInfo[1, 2:ncol(metaInfo)]
  getMask = as.matrix(metaInfo[2, 2:ncol(metaInfo)])
  
  # validate schema 
  X = dropInvalidType(F, getSchema)

  if(sum(getMask) > 0 )
  {
    # always recode the label
    index = vectorToCsv(getMask)
    jspecR = "{ids:true, recode:["+index+"]}"
    [eX, X_meta] = transformencode(target=X, spec=jspecR);
    # change the schema
    getSchema = map(getSchema, "x->x.replace(\"STRING\", \"INT64\")")
    getSchema = map(getSchema, "x->x.replace(\"BOOLEAN\", \"INT64\")")
    
  } else
    eX = as.matrix(X)

  eY = eX[, ncol(eX)]
  eX = eX[, 1:ncol(X) - 1]
  Xout = outlier(eX, 0.25)
    
  getMask = getMask[, 1:ncol(getMask) - 1] # strip the mask of class label
  getSchema = getSchema[, 1:ncol(getSchema) - 1] # strip the mask of class label

  
  [eX, eY] = doSample(eX, eY, sample)
  # get train test and validation set with balanced class distribution
  [X_train, y_train, X_test, y_test] = splitBalanced(eX, eY, 0.8, FALSE)

  params = list("reg", "maxi");
  paramRanges = list(10^seq(0,-10), seq(10,100, 10));

  [opt, loss] = gridSearchMLR(X_train, y_train, X_test, y_test, 
    "multiLogReg", "lossFunc", params, paramRanges, FALSE);
   
  d_accuracy = classifyDirty(X_train, y_train, opt, getMask, isWeighted, cv)
  # [eX, eY] = prioritise(eX, eY, getMask)



  [pip, hp, acc] = bandit(X_train=X_train, Y_train=y_train,  mask=getMask, MLhp=opt,
    schema=getSchema, lp=logical, primitives=primitives, param=param, k=k, testAccuracy=d_accuracy,
    isWeighted=isWeighted, R=R, cv=cv, verbose=TRUE);
  # acc = as.matrix(1)
  # pip = frame(["imputeByMean", "winsorize", "imputeByMean", "scale", "dummycoding", "pca"], rows=1, cols=6)
  # print("pipeline "+toString(pip))
  # hp = matrix("4.000 1.000 0.000 0.000 2.000 " 
              # +"4.000 0.000 0.000 1.000 0.000 " 
              # +"4.000 1.000 0.000 0.000 2.000 " 
              # +"6.000 1.000 0.000 0.000 0.000 0.000 0.000 " 
             # + "4.000 1.000 0.000 0.000 2.000 " 
              # +"7.000 2.000 1.000 0.000 0.000 0.000 0.000 2.000 0.000 0.000 0.000 0.000 0.000", 

  tmp_hp = cbind(matrix(NaN, nrow(hp), 1), hp)
  result = cbind(pip, as.frame(tmp_hp))
  result = cbind(result, as.frame(acc))
  
  if(as.scalar((is.na(acc[1,1]))) == 1 | as.scalar(acc[1,1]) < d_accuracy)
    stop("warning: no best pipeline found")
  
  clean_accuracy = testBestPipeline(pip=pip[1,], hp=hp[1,], X_train=X_train, y_train=y_train,
    X_test=X_test, y_test=y_test, cmask=getMask, MLhp=opt, schema=getSchema, valAcc=as.scalar(acc[1,1]), dirAcc=d_accuracy,
    isWeighted=isWeighted)
    
  res = cbind(as.matrix(d_accuracy), as.matrix(clean_accuracy))
}

testBestPipeline = function(Frame[Unknown] pip, Matrix[Double] hp, Matrix[Double] X_train, Matrix[Double] y_train, 
  Matrix[Double] X_test, Matrix[Double] y_test, Matrix[Double] cmask, Matrix[Double] MLhp, Frame[Unknown] schema,
  Double valAcc, Double dirAcc, Boolean isWeighted)
  return (Double result) {
  print("hp "+toString(hp))
  ls = list();
  i = 1; k = 1
  trRow=nrow(X_train)
  # construct the parameter list for best hyper-parameters
  while(k <= ncol(pip))
  {
    end = as.integer(i+as.integer(as.scalar(hp[1,i])))
    mat = hp[1, i+1:end]
    i = end + 1
    ls = append(ls, mat)
    k = k + 1
  }
  
  # clean using best pipeline and train model
  [X_train, Y_train] = executePipeline(pip, rbind(X_train,X_test),
    rbind(y_train,y_test), cmask, ls, FALSE)
  X_train_clean = X_train[1:trRow, ]
  y_train_clean = Y_train[1:trRow, ]
  X_test_clean = X_train[trRow+1:nrow(X_train), ]
  y_test_clean = Y_train[trRow+1:nrow(X_train), ]

  # classify after cleaning  
  betas = multiLogReg(X=X_train_clean, Y=y_train_clean, icpt=1,
    reg=as.scalar(MLhp[1,1]), tol= 1e-9, maxi=as.scalar(MLhp[1,2]), 
    maxii= 50, verbose=FALSE);
    
  [c_prob, c_yhat, c_accuracy] = multiLogRegPredict(X_test_clean, betas, y_test_clean, FALSE)
    c_accuracy = getAccuracy(y_test_clean, c_yhat, isWeighted)
  [confusionCount_c, confusionAVG_c] = confusionMatrix(P=c_yhat, Y=y_test_clean)
  
  print("accuracy of dirty data  "+dirAcc)
  print("accuracy of val data  "+valAcc)
  print("accuracy of test accuracy "+c_accuracy)
  print("clean confusion matrix  \n"+toString(confusionCount_c))
  
  result = c_accuracy
}

# stratified sampling 
doSample = function(Matrix[Double] eX, Matrix[Double] eY, Double ratio)
  return (Matrix[Double] eX, Matrix[Double] eY)
{
  MIN_SAMPLE = 10000
  sampled = floor(nrow(eX) * ratio)
  sample = ifelse(sampled > MIN_SAMPLE, TRUE, FALSE)
  if(sample)
  {
    XY = order(target = cbind(eY, eX),  by = 1, decreasing=FALSE, index.return=FALSE)
    # get the class count 
    classes = table(eY, 1)
    print("classes")
    print(toString(classes))
    while(FALSE){}
    start_class = 1
    out_s = 1 
    out_e = 0
    end_class = 0
  
    out = matrix(0, sampled, ncol(XY))
    classes_ratio = floor(classes*ratio)
    print("class ratio "+toString(classes_ratio))
    for(i in 1:nrow(classes))
    {
      end_class = end_class + as.scalar(classes[i])
      class_t = XY[start_class:end_class, ]
      out_e = out_e + as.scalar(classes_ratio[i]) 
      out[out_s:out_e, ] = class_t[1:as.scalar(classes_ratio[i]), ] 
      out_s = out_e + 1
      start_class = end_class + 1
    }
    out = removeEmpty(target = out, margin = "rows")
    eY = out[, 1]
    eX = out[, 2:ncol(out)]
  }
}

# incomplete implementation of automatic logical pipelines
generateLogical = function(Matrix[Double] X, Matrix[Double] Y, Matrix[Double] mask)
{
  
  # detection = 
  logical = as.frame("")
  no_of_mv = sum(is.na(X))
  X = replace(target= X, pattern = NaN, replacement = 0)
  
  # get the stats
  colMin = colMins(X)
  colMax = colMaxs(X)
  colMean = colMeans(X)
  colSd = colSds(X)
  count3sdplus = sum(X > (colMean + 3*colSd )) 
  count3sdminus = sum(X < (colMean - 3*colSd )) 
  outliers = count3sdplus + count3sdminus
  ctab = table(Y, 1)
  minCatPer = min(ctab) / nrow(ctab)
  maxCat = max(ctab) / nrow(ctab)
  
  mv_to_data_ratio = no_of_mv/(nrow(X) * ncol(X))
  out_to_data_ratio = outliers/ (nrow(X) * ncol(X))
  
  if(mv_to_data_ratio > 0.1)
    logical = cbind(logical, as.frame("MVI"))
  if(out_to_data_ratio > 0.1)
    logical = cbind(logical, as.frame("OTLR"))
  if((maxCat - minCatPer) > 0.3)
    logical = cbind(logical, as.frame("CI"))
  # if(min(X))
}

# stratified splitting 
getDataSplits = function(Matrix[Double] X, Matrix[Double] Y, Double splitRatio, Boolean verbose)
return (Matrix[Double] X_train, Matrix[Double] y_train, Matrix[Double] X_test, 
        Matrix[Double] y_test) 
{


  XY = order(target = cbind(Y, X),  by = 1, decreasing=FALSE, index.return=FALSE)
  # get the class count 
  classes = table(Y, 1)

  split = floor(nrow(X) * splitRatio)
  start_class = 1
  train_row_s = 1 
  test_row_s = 1 
  train_row_e = 0
  test_row_e = 0
  end_class = 0
  
  outTrain = matrix(0, split+nrow(classes), ncol(XY))
  outTest =  matrix(0, (nrow(X) - split)+nrow(classes), ncol(XY))
  
  classes_ratio_train = floor(classes*splitRatio)
  classes_ratio_test = classes - classes_ratio_train
  if(verbose) {
    print("rows "+nrow(X))
    print("classes \n"+toString(classes))
    print("train ratio \n"+toString(classes_ratio_train))
    print("test ratio \n"+toString(classes_ratio_test))
  }
  for(i in 1:nrow(classes))
  {
    end_class = end_class + as.scalar(classes[i])
    class_t = XY[start_class:end_class, ]

    train_row_e = train_row_e + as.scalar(classes_ratio_train[i]) 
    test_row_e = test_row_e + as.scalar(classes_ratio_test[i]) 
    outTrain[train_row_s:train_row_e, ] = class_t[1:as.scalar(classes_ratio_train[i]), ]   
    outTest[test_row_s:test_row_e, ] = class_t[as.scalar(classes_ratio_train[i])+1:nrow(class_t), ]

    train_row_s = train_row_e + 1
    test_row_s = test_row_e + 1
    start_class = end_class + 1
  }

  outTrain = removeEmpty(target = outTrain, margin = "rows")
  outTest = removeEmpty(target = outTest, margin = "rows")
  y_train = outTrain[, 1]
  X_train = outTrain[, 2:ncol(outTrain)]
  y_test = outTest[, 1]
  X_test = outTest[, 2:ncol(outTest)]
}

# entropy calculation for finding class imbalance
getBalanceScore = function(Matrix[Double] Y)
return (Matrix[Double] isBalanced)
{
  # get count of instances in each class
  k  = table(Y, 1)
  n = nrow(Y)
  # compute Shannon entropy for i to k H = − ∑ ci/n * log(ci/n)
  H = -sum((k/n) * log(k/n))
  # Balance = H / log(k) return 0 for unbalance data and 1 for balanced data
  isBalanced = H / log(k)
}


prioritise = function(Matrix[Double] X, Matrix[Double] y, Matrix[Double] mask)
return(Matrix[Double] sortedX, Matrix[Double] sortedY)
{
# classify without cleaning fill with the default values 1
  Xtrain = replace(target = X, pattern = NaN, replacement=1)
  
  train_max_values = colMaxs(Xtrain)
  dX_train = matrix(0,nrow(Xtrain),0)
  # one-hot-encode the categorical features
  for(i in 1:ncol(mask))
  {
    if(as.scalar(mask[1, i]) == 1)
    {
      max_value = as.scalar(train_max_values[1, i])
      while(FALSE){}
      encoded = toOneHot(Xtrain[, i], max_value)
      dX_train = cbind(dX_train, encoded)
    }
    else {
      dX_train = cbind(dX_train, Xtrain[, i])
    }

  }
  # print('ncol in dx train '+ncol(dX_train))
  while(FALSE){}
  # learn model
  B = multiLogReg(X=dX_train, Y=y, icpt=1, reg=0, tol=1e-6, maxi = 100, maxii= 0, verbose=FALSE);
  [prob,yhat,accuracy] = multiLogRegPredict(X=dX_train, B=B, Y=y, verbose=FALSE);
  # print("the accuracy "+accuracy)
  
  detect = (rowMaxs(prob) < 0.9)
  position = seq(1, nrow(X))
  err = removeEmpty(target=position, margin="rows", select=detect)
  # print("erroneous records position "+toString(err))

  position = seq(1, nrow(X))
  fn = removeEmpty(target=position, margin="rows", select=(detect==0))
  # print("fine records position "+toString(fn))
  
  erroneousX = removeEmpty(target=X, margin="rows", select=detect)
  erroneousY = removeEmpty(target=y, margin="rows", select=detect)
  # print("errorneous X \n"+toString(cbind(erroneousX, erroneousY), rows=200))

  fineX = removeEmpty(target=X, margin="rows", select=(detect == 0))
  fineY = removeEmpty(target=y, margin="rows", select=(detect == 0))
  # print("error waley "+toString(erroneous))
  # print("fine X \n"+toString(cbind(fineX, fineY), rows=500))
  sortedX = rbind(erroneousX, fineX)
  sortedY = rbind(erroneousY, fineY)
  sortedX = removeEmpty(target=sortedX, margin="rows", select=sortedY)
  sortedY = removeEmpty(target=sortedY, margin="rows", select=sortedY)
 
  # print("YES "+(nrow(X) == nrow(sortedX)))
  
}

dummycoding = function(Matrix[Double] X, Matrix[Double] mask)
return (Matrix[Double] dX_train) {

  idx = vectorToCsv(mask)
  
  # specifications for one-hot encoding of categorical features
  jspecDC = "{ids:true, dummycode:["+idx+"]}";
  # OHE of categorical features
  [dX_train, dM] = transformencode(target=as.frame(X), spec=jspecDC);

  
  # train_max_values = colMaxs(Xtrain)
  # dX_train = matrix(0,nrow(Xtrain),0)
  # # one-hot-encode the categorical features
  # for(i in 1:ncol(mask))
  # {
    # if(as.scalar(mask[1, i]) == 1)
    # {
      # max_value = max(train_max_values[1, i])
      # encoded = toOneHot(Xtrain[, i], max_value)
      # dX_train = cbind(dX_train, encoded)
    # }
    # else {
      # dX_train = cbind(dX_train, Xtrain[, i])
    # }

  # }
}


gridSearchMLR = function(Matrix[Double] Xtrain, Matrix[Double] ytrain, Matrix[Double] Xtest,
    Matrix[Double] ytest, String train, String predict,	List[String] params, List[Unknown] paramValues,
    Boolean verbose = TRUE) 

  return (Matrix[Double] opt, Matrix[Double] Rloss) 
{
  # Step 0) preparation of parameters, lengths, and values in convenient form
  numParams = length(params);
  paramLens = matrix(0, numParams, 1);
  for( j in 1:numParams ) {
    vect = as.matrix(paramValues[j,1]);
    paramLens[j,1] = nrow(vect);
  }
  paramVals = matrix(0, numParams, max(paramLens));
  for( j in 1:numParams ) {
    vect = as.matrix(paramValues[j,1]);
    paramVals[j,1:nrow(vect)] = t(vect);
  }
  cumLens = rev(cumprod(rev(paramLens))/rev(paramLens));
  numConfigs = prod(paramLens);
  
  # Step 1) materialize hyper-parameter combinations 
  # (simplify debugging and compared to compute negligible)
  HP = matrix(0, numConfigs, numParams);
  parfor( i in 1:nrow(HP) ) {
    for( j in 1:numParams )
      HP[i,j] = paramVals[j,as.scalar(((i-1)/cumLens[j,1])%%paramLens[j,1]+1)];
  }

  if( verbose )
    print("GridSeach: Hyper-parameter combinations: \n"+toString(HP));

  # Step 2) training/scoring of parameter combinations
  # TODO integrate cross validation
 
  Rloss = matrix(0, nrow(HP), 3);
  arguments1 = list(X=Xtrain, Y=ytrain, icpt=1, reg=-1, tol=1e-9, maxi=-1, maxii=0, verbose=FALSE);

  parfor( i in 1:nrow(HP)) {
    # a) replace training arguments
    largs1 = arguments1;

    for( j in 1:numParams ) {
      largs1[as.scalar(params[j])] = as.scalar(HP[i,j]);
    }
    # b) core training/scoring and write-back
    # TODO investigate rmvar handling with explicit binding (lbeta)
    Rbeta1 = eval(train, largs1);
    Rloss[i,1] = eval(predict, list(Xtest, ytest, Rbeta1));
  }

  # Step 3) select best parameter combination
  ix = as.scalar(rowIndexMin(t(Rloss[,2])));
  opt = HP[ix,]; # optimal hyper-parameters
 
}


lossFunc = function(Matrix[Double] X, Matrix[Double] y, Matrix[Double] B) 
return (Matrix[Double] loss) {
  [prob, yhat, acc] = multiLogRegPredict(X=X, B=B, Y=y,  verbose=FALSE)
  loss = as.matrix(1 - (acc/100))
  [confusionCount_c, confusionAVG_c] = confusionMatrix(P=yhat, Y=y)
  # print(" loss: "+toString(loss))
  # print("confusion matrix")
  # print(toString(confusionCount_c))
}


classifyDirty = function(Matrix[Double] Xtrain, Matrix[Double] ytrain, Matrix[Double] opt, 
  Matrix[Double] mask, Boolean isWeighted = TRUE, Integer cv)
  return (Double accuracy)
{
  # classify without cleaning fill with edfault values 1
  Xtrain = replace(target = Xtrain, pattern = NaN, replacement=1)
  
  dX_train = dummycoding(Xtrain, mask)

  accuracy = crossV(Xtrain, ytrain, cv, mask, opt, isWeighted)
  accuracy = mean(accuracy)

  # # learn model
  # B = multiLogReg(X=dX_train, Y=ytrain, icpt=2, reg=as.scalar(opt[1,1]), maxi = as.scalar(opt[1,2]), maxii= 0, verbose=FALSE);
  # [M,pred,accuracy] = multiLogRegPredict(X=dX_test, B=B, Y=ytest, verbose=FALSE);

  # if(isWeighted) 
    # accuracy = getAccuracy(y=ytest, yhat=pred, isWeighted=isWeighted)
  print("cross validated dirty accuracy "+accuracy)
}


crossV = function(Matrix[double] X, Matrix[double] y, Integer k, Matrix[Double] mask,
  Matrix[Double] MLhp, Boolean isWeighted) 
return (Matrix[Double] accuracyMatrix)
{
  accuracyMatrix = matrix(0, k, 1)

  dataList = list()
  testL = list()
  data = order(target = cbind(y, X),  by = 1, decreasing=FALSE, index.return=FALSE)
  classes = table(data[, 1], 1)
  ins_per_fold = classes/k
  start_fold = matrix(1, rows=nrow(ins_per_fold), cols=1)
  fold_idxes = cbind(start_fold, ins_per_fold)

  start_i = 0; end_i = 0; idx_fold = 1;;
  for(i in 1:k)
  {
    fold_i = matrix(0, 0, ncol(data))
    start=0; end=0; 
    for(j in 1:nrow(classes))
    {
      idx = as.scalar(classes[j, 1])
      start = end + 1;
      end = end + idx
      class_j =  data[start:end, ]


      start_i = as.scalar(fold_idxes[j, 1]);
      end_i = as.scalar(fold_idxes[j, 2])

      fold_i = rbind(fold_i, class_j[start_i:end_i, ])
    }

    dataList = append(dataList, fold_i)
    fold_idxes[, 1] = fold_idxes[, 2] + 1
    fold_idxes[, 2] += ins_per_fold
  }

  for(i in seq(1,k))
  {
      [trainList, hold_out] = remove(dataList, i)
      trainset = rbind(trainList)
      testset = as.matrix(hold_out)
      trainX = trainset[, 2:ncol(trainset)]
      trainy = trainset[, 1]
      testX = testset[, 2:ncol(testset)]
      testy = testset[, 1]
      beta = multiLogReg(X=trainX, Y=trainy, icpt=1, reg=as.scalar(MLhp[1,1]), tol= 1e-9, 
      maxi=as.scalar(MLhp[1,2]), maxii= 50, verbose=FALSE);
      [prob, yhat, a] = multiLogRegPredict(testX, beta, testy, FALSE)
      accuracy = getAccuracy(testy, yhat, isWeighted)
      accuracyMatrix[i] = accuracy
  }

}
