#-------------------------------------------------------------
#
# 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.
#
#-------------------------------------------------------------

# This function computes the measure of impurity for the given dataset based on the passed method (gini or entropy).
# The current version expects the target vector to contain only 0 or 1 values.
#
# INPUT:
# --------------------------------------------------------------------------
# X        Feature matrix.
# Y        Target vector containing 0 and 1 values.
# R        Vector indicating whether a feature is categorical or continuous.
#          1 denotes a continuous feature, 2 denotes a categorical feature.
# n_bins   Number of bins for binning in case of scale features.
# method   String indicating the method to use; either "entropy" or "gini".
# --------------------------------------------------------------------------
#
# OUTPUT:
# ------------------------------------------------------------------------
# IM     (1 x ncol(X)) row vector containing information/gini gain for
#        each feature of the dataset.
#        In case of gini, the values denote the gini gains, i.e. how much
#        impurity was removed with the respective split. The higher the
#        value, the better the split.
#        In case of entropy, the values denote the information gain, i.e.
#        how much entropy was removed. The higher the information gain,
#        the better the split.
# ------------------------------------------------------------------------

m_impurityMeasures = function(Matrix[Double] X, Matrix[Double] Y, Matrix[Double] R, Integer n_bins = 20, String method)
  return (Matrix[Double] IM)
{
  if (method != "entropy" & method != "gini") {
    stop("Please specify the correct method - should be either entropy or gini.")
  }

  IM = matrix(0.0, rows = 1, cols = ncol(X))

  parfor (i in 1:ncol(X)) {
    if (as.scalar(R[,i]) == 1) {
      binned_feature = applyBinning(X[,i], n_bins)
      IM[,i] = getImpurityMeasure(binned_feature, Y, n_bins, method)
    } else {
      IM[,i] = getImpurityMeasure(X[,i], Y, max(X[,i]), method)
    }
  }
}

getImpurityMeasure = function(Matrix[Double] feature, Matrix[Double] Y, Double max_cat, String method)
  return (Double gain)
{
  n_true_labels = sum(Y)
  n_false_labels = length(Y) - n_true_labels
  parent_impurity = calcImpurity(n_true_labels, n_false_labels, length(feature), method)

  # calculate the impurity after the split
  children_impurity = 0
  for (i in 1:max_cat) {
    count_true = 0
    count_false = 0
    for (j in 1:length(feature)) {
      if (as.scalar(feature[j,]) == i) {
        if (as.scalar(Y[j,]) == 0) {
          count_false += 1
        } else {
          count_true += 1
        }
      }
    }
    if (!(count_true == 0 & count_false == 0)) {
      children_impurity = children_impurity + calcImpurity(count_true, count_false, length(feature), method)
    }
  }
  gain = parent_impurity - children_impurity
}

calcImpurity = function(Double n_true, Double n_false, Double n_vars, String method)
  return (Double impurity)
{
  impurity = 0
  prob_true = n_true / (n_true + n_false)
  prob_false = n_false / (n_true + n_false)
  weight = (n_true + n_false) / n_vars

  if (prob_true != 1 & prob_false != 1) { # if there is more than one class, calculate new impurity according to method.
    if (method == "entropy") { # dividing by log(2) to obtain the information gain in bits
      impurity = (-1) * weight * (prob_true * log(prob_true)/log(2) + prob_false * log(prob_false)/log(2))
    } else if (method == "gini") {
      impurity = weight * (1 - (prob_true^2 + prob_false^2))
    }
  }
}

applyBinning = function(Matrix[Double] feature, Double n_bins)
  return (Matrix[Double] output_f)
{
  # equi-width binning.

  if (length(feature) < n_bins) {
    n_bins = length(feature)
  }
  max_v = max(feature)
  min_v = min(feature)
  width = (max_v - min_v) / n_bins
  output_f = matrix(1, rows = nrow(feature), cols = 1)

  parfor (i in 1:length(feature)) {
    binned = FALSE
    j = 1
    while (binned == FALSE) {
      if (as.scalar(feature[i,]) <= min_v + j * width) {
        output_f[i,] = j
        binned = TRUE
      }
      j += 1
    }
  }
}
