#-------------------------------------------------------------
#
# 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.
#
#-------------------------------------------------------------
# Generate the logical pipelines using basic evolutionary algorithm, 
# population -> logical Pipeline, chromosome -> physical pipeline, gene -> hp
# input :
# 1. Dataset X
# 2. population, different logical seed pipelines format = [number of operators, op1, op2, op3 ..., opn]
# 3. number of iterations
# 4. pipLength, length of pipeline, number of operator in each pipeline
# 5. meta data list i.e., schema, mask, fdmask
# 6. target list i.e, target application, cv value etc.
# 7. primitives, physical operator list
# 8. param, physical operator hyperparameters
# 9. num_inst value, number of physical instances for each logical to be executed
# 10. num_exec, how many times each physical pipeline should be executed
# 11. n_pop, children created in each generation
# output: best logical pipeline and evaluation time in ms

# idea is to get the initial set of logical pipelines, as population, then get the num_inst physical pipelines foreach
# logical pipeline in population. Then execute these physical pipelines num_exec time were in each execution a random
# set of hyperparameters is used to execute operators. 
# The compute a score vector by storing the best score foreach logical pipeline in population. Sort the pipelines by
# score and take n_pop pipelines as parents for generating new population from the selected pipelines take a pair in
# each iteration as parent and generate a pair of children by doing crossover and mutation.
# There are total 3 transformation that will be performed to create a new chromosomes (children)
#   1. crossover to create children by adding n operation from p1 to p2 and vice versa. 
#   2. mutation to swap two operations in the children based on a mutation rate, swap is randomly based on mutationRatio.
#   3. removal to remove n operations from a child
#   4. Add to randomly add a new operation in existing child
# These new children will be the population in next iteration.
# Repeat the process max_iter time. Converge in between if the best_score of previous generation is better then
# best_score of new generation.


source("scripts/builtin/bandit.dml") as bandit;

enumerateLogical = function(Matrix[Double] X, Matrix[Double] y, Matrix[Double] Xtest, Matrix[Double] ytest,
  Frame[Unknown] initial_population, Frame[String] refSol = as.frame("NaN"), Integer seed = -1, Integer max_iter=10,
  List[Unknown] metaList, String evaluationFunc, Matrix[Double] evalFunHp,
  Frame[Unknown] primitives, Frame[Unknown] param, Double dirtyScore = 79, Boolean cv=FALSE, Boolean cvk=3,
  Boolean verbose, List[Unknown] ctx=list(prefix="----"))
return (Frame[Unknown] outputPip, Matrix[Double] outputHp, boolean converged, Double refChanges, Frame[Unknown] acc)
{

  # # # initialize mask and other meta variables
  mask = as.matrix(metaList['mask'])
  prefix = as.scalar(ctx["prefix"]);  
  iter = 1

  # # # initialize variables for convergence check. If there is not improvement in last three iterations then we converge
  convCheck = matrix(0, rows=3, cols=1)
  convCounter = 0
  converged = FALSE
  
  # # # get the operations to add in transitions
  [allOps, ref] = getOps(param[, 2], refSol, as.scalar(metaList['distY']), nrow(y), min(y))

  # unrolled by physical pipelines; get the physical pipeline
  pipelines = frame(0, rows=nrow(primitives)^ncol(primitives), cols=max(ncol(initial_population), ncol(ref)))
  start = 1; 
  end = 0;
  for(i in 1:nrow(initial_population)) { 
    pconf = bandit::get_physical_configurations(initial_population[i], 0, primitives)
    end = end + nrow(pconf)
    pipelines[start:end, 1:ncol(pconf)] = pconf
    start = end + 1
  }
  pipelines = removeEmpty(target = pipelines, margin="rows") 
  
  # # # initialize output variable to store final pipelines and hyp
  FLAGS = 5
  ROWS = (max_iter+1)*(nrow(pipelines))
  finalOutputFrame = frame(0, rows=ROWS, cols=max_iter*2)
  # num of max operations * max hp per operation * no of flag + buffer for pipeline no and acc
  maxParam = ncol(finalOutputFrame) * max(as.matrix(param[, 3])) * FLAGS + 2 
  finalOutputMatrix = matrix(0, rows=ROWS, cols=maxParam)
  
  # # if the data has categorical columns then add the dummycode operation
  if(sum(mask) > 0)
  {
    dummyEncode = frame("dummycoding", rows=nrow(pipelines), cols=1)
    pipelines[, 2] = dummyEncode
  }
  
  pipelines = rbind(ref, pipelines)
  # # # treat the pipelines as initial population
  population = pipelines
  populationSize = nrow(pipelines)
  
  # # initialize the transitions add, remove, mutate and crossover
  transitions = sample(4, (populationSize * max_iter), TRUE, seed)
  opToAdd = sample(nrow(allOps), (populationSize * max_iter), TRUE, seed)
  
  # # # initialize the indexing variables to store outputs of all iterations
  outputFrameStart = 1
  outputFrameEnd = 0
  while(!converged & iter <= max_iter)
  {
    # populationLength = max(populationLength, ncol(population))
    id = matrix(seq(1, nrow(population)*2), rows=nrow(population), cols=2)
    print(prefix+" EnumLP iteration "+iter+"/"+as.integer(max_iter)+":" );
    
    # # # execute the physical pipelines
    [outPip, outHp, p, refChanges] = bandit::run_with_hyperparam(ph_pip=cbind(as.frame(id), population),
      X=X, Y=y, Xtest=Xtest, Ytest=ytest, metaList=metaList, evaluationFunc=evaluationFunc, evalFunHp=evalFunHp, param=param,
      cv=cv, cvk=cvk, seed=seed, default=TRUE)
    
    # # sort the configurations score-wise
    actPip = cbind(as.frame(outPip[, 1]), as.frame(refChanges))
    actPip = cbind(actPip, population)
    sort_mask = cbind(matrix(0, rows=1, cols=2), matrix(1, rows=1, cols=ncol(population)))
    sortedPipelines = frameSort(actPip, sort_mask, TRUE)
    sortedHp = order(target = outHp, by = 1, decreasing=TRUE)
    
    # # fix the convergence: converge when there is no change in the last three iterations
    bestSoFar = as.double(as.scalar(sortedPipelines[1, 1]))
    if(bestSoFar > min(convCheck))
    {
      idx = as.scalar(rowIndexMin(t(convCheck)))
      convCheck[idx] = bestSoFar
    }
    else convCounter =  convCounter + 1
    converged = convCounter > nrow(convCheck) #as.double(as.scalar(sortedPipelines[1, 1])) > dirtyScore

    # # if not converged then keep the top-k successive halving
    if(converged)
      print(prefix+" EnumLP  converged after "+iter+" / "+max_iter+" iterations")  
    diR = round(nrow(sortedPipelines)/2)
    if(nrow(sortedPipelines) > 1) {
      sortedPipelines = sortedPipelines[1:diR]
      sortedHp = sortedHp[1:diR]
    }
    
    # # # store the topk in final output variables
    outputFrameEnd = outputFrameEnd + nrow(sortedPipelines) 
    finalOutputFrame[outputFrameStart:outputFrameEnd, 1:ncol(sortedPipelines)] = sortedPipelines
    finalOutputMatrix[outputFrameStart:outputFrameEnd, 1:ncol(sortedHp)] = sortedHp
    outputFrameStart = outputFrameEnd + 1

    # # # if converged then stop otherwise generate new population
    children = frame(0, rows=populationSize, cols=ncol(sortedPipelines)+(ncol(sortedPipelines)/2))
    sortedPipelines = sortedPipelines[, 3:ncol(sortedPipelines)]
    topk = finalOutputFrame
    sort_mask = cbind(matrix(0, rows=1, cols=2), matrix(1, rows=1, cols=ncol(topk) - 2))
    topk = removeEmpty(target=topk, margin="rows")
    topk = frameSort(topk, sort_mask, TRUE)
    topk = topk[, 3:ncol(topk)]
    
    # # randomly pick the pipelines for transitions
    pipRand = sample(nrow(sortedPipelines), populationSize, TRUE, seed)
    if(!converged) {
      parfor(i in 1:nrow(children), check=0) {
        idxR = (nrow(children) * (iter - 1)) + i
        idx = as.scalar(pipRand[i])
        top = removeEmpty(target=topk[idx], margin="cols")
        idx2 = min(max(pipRand), idx + 1)
        top2 = removeEmpty(target=topk[idx2], margin="cols")
        # # # keep the tail "dummycode" operation from transitions
        if(sum(mask) > 0) {
          tail = top[, ncol(top)]
          tail2 = top2[, ncol(top2)]
          top = top[, 1:ncol(top) - 1]
          top2 = top2[, 1:ncol(top2) - 1]
        }
          
        random = ifelse(ncol(top) <=2, 1, as.scalar(transitions[idxR]))
        if(random == 1)
          c1 = addition(top, allOps[as.scalar(opToAdd[idxR])]) 
        else if(random == 2)
          c1 = mutation(top, seed) 
        else if(random == 3)
          c1 = removal(top, seed) 
        else if(random == 4)
          c1 = crossover(top, top2, seed)
        
        # # # put back the tail operation 
        if(sum(mask) > 0)
          c1 = cbind(c1, tail)
        children[i, 1:ncol(c1)] = c1
      }
    }
    population = removeEmpty(target=children, margin="cols")
    iter  = iter + 1
  }
  if(!converged) {
    print(prefix+" EnumLP did not converge after "+(iter - 1)+" / "+max_iter+" iterations")  
  }

  # # sort the final output in increasing order
  outputPip = finalOutputFrame
  sort_mask = cbind(matrix(0, rows=1, cols=2), matrix(1, rows=1, cols=ncol(outputPip) - 2))
  outputPip = removeEmpty(target=outputPip, margin="rows")
  outputPip = frameSort(outputPip, sort_mask, FALSE)
  
  refChanges = as.double(as.scalar(outputPip[nrow(outputPip), 2]))
  acc = outputPip[, 1]

  outputPip = outputPip[,3:ncol(outputPip)]
  # # # prepare the hyp output
  hpLength = ((ncol(outputPip) + 2) * FLAGS * 3) + 1 
  outputHp = finalOutputMatrix[, 1:hpLength]
  outputHp = order(target = outputHp, by = 1, decreasing=FALSE)
}

addition = function(Frame[Unknown] top, Frame[Unknown] opToAdd)
return (Frame[Unknown] child)
{
  # # # never add same operation adjacent to each other
  if(as.scalar(top[1,1]) != as.scalar(opToAdd[1,1]))
    child = cbind(opToAdd, top)
  else 
    child = cbind(top, opToAdd)
}

mutation = function(Frame[Unknown] child, Integer seed)
return (Frame [Unknown] mChild)
{
  if(ncol(child) >= 2)
  {
    r = sample(ncol(child), 2, FALSE, seed)
    r1 = as.scalar(r[1,1])
    r2 = as.scalar(r[2,1])
    temp = child[1, r1]
    child[1, r1] = child[1, r2]
    child[1, r2] = temp
  }
  mChild = child
}

removal = function(Frame[Unknown] child, Integer seed)
return (Frame[Unknown] child)
{
  if(ncol(child) >= 2)
  {
    pos = as.scalar(sample(ncol(child), 1, FALSE, seed))
    child[1, pos] = as.frame(0)
    child = removeEmpty(target=child, margin="cols")
  }
}

crossover = function(Frame[Unknown] p1, Frame[Unknown] p2, Integer seed)
return(Frame[Unknown] child)
{
    # # randomly select the lengths to be append
    lp1 = as.scalar(sample(ncol(p1), 1, FALSE, seed))
    lp2 = as.scalar(sample(ncol(p2), 1, FALSE, seed))
    child = cbind(p1[, 1:lp1], p2[, lp2:ncol(p2)])
}


getOps = function( Frame[string] allOps, Frame[String] refSol, Integer dist, Integer n, Integer minValue)
 return (Frame[String] allOps, Frame[String] refSol) {
 
  # # # TODO fix the following hard-coded condition by taking a file input
  # # allOps are the operation which are randomly added to a population, for now I am reusing param file
  # # so map condition with remove the operations which should not be added twice in a pipeline i.e., dummycoding
  # # for regression class imbalance operators are also removed
  if(n > 0 & minValue >= 1 & dist <= 15) {
    allOps = map(allOps, "x -> (!x.equals(\"dummycoding\") & !x.equals(\"frequencyEncode\")
    & !x.equals(\"dbscan\") & !x.equals(\"WoE\") & !x.equals(\"pca\") & !x.equals(\"ppca\") &
     !x.equals(\"abstain\") & !x.equals(\"underSampling\") & !x.equals(\"flipLabels\") & !x.equals(\"mice\") & !x.equals(\"SMOTE\"))?x:\"0\"")
    ref = frame(["imputeByMean", "winsorize", "scale", "dummycoding"], rows=1, cols=4)
  }
  else {
    allOps = map(allOps, "x -> (!x.equals(\"dummycoding\") & !x.equals(\"frequencyEncode\") & !x.equals(\"tomeklink\")
      & !x.equals(\"dbscan\") & !x.equals(\"WoE\") & !x.equals(\"pca\") & !x.equals(\"ppca\") &
      !x.equals(\"abstain\") & !x.equals(\"underSampling\") & !x.equals(\"flipLabels\") & !x.equals(\"mice\") & !x.equals(\"SMOTE\"))?x:\"0\"") 
    ref = frame(["imputeByMean", "winsorize", "scale", "dummycoding"], rows=1, cols=4)
  }
  if(as.scalar(refSol[1,1]) == "NaN")
    refSol = ref
  allOps = removeEmpty(target=allOps, margin="rows")
}