#-------------------------------------------------------------
#
# 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 builtin function implements a NaiveBayes classification.
#
# INPUT:
# ------------------------------------------------------------------------------
# D        Input feature matrix of shape N x M
# C        Class label vector (positive integers) of shape N x 1.
# laplace  Laplace smoothing correction (prevent zero probabilities)
# verbose  Flag for verbose debug output
# ------------------------------------------------------------------------------
#
# OUTPUT:
# ------------------------------------------------------------------------------
# prior              Class prior probabilities
# classConditionals  Class conditional feature distributions
# ------------------------------------------------------------------------------

m_naiveBayes = function(Matrix[Double] D,
  Matrix[Double] C, Double laplace = 1, Boolean verbose = TRUE)
  return (Matrix[Double] prior, Matrix[Double] classConditionals)
{
  laplaceCorrection = laplace;
  numRows = nrow(D)
  numFeatures = ncol(D)
  minFeatureVal = min(D)
  numClasses = as.integer(max(C))
  minLabelVal = min(C)

  # sanity checks of data and arguments
  if(minFeatureVal < 0)
    stop("naiveBayes: Stopping due to invalid argument: Multinomial naive Bayes "
       + " is meant for count-based feature values, minimum value in D is negative")
  if(numRows < 2)
    stop("naiveBayes: Stopping due to invalid inputs: "
       + "Not possible to learn a classifier without at least 2 rows")
  if(minLabelVal < 1)
    stop("naiveBayes: Stopping due to invalid argument: Label vector (C) must be recoded")
  if(numClasses == 1)
    stop("naiveBayes: Stopping due to invalid argument: "
       + "Maximum label value is 1, need more than one class to learn a multi-class classifier")
  if(sum(abs(C%%1 == 0)) != numRows)
    stop("naiveBayes: Stopping due to invalid argument: "
       + "Please ensure that C contains (positive) integral labels")
  if(laplaceCorrection < 0)
    stop("naiveBayes: Stopping due to invalid argument: "
       + "Laplacian correction (laplace) must be non-negative")

  # Compute conditionals
  # Compute the feature counts for each class
  classFeatureCounts = aggregate(target=D, groups=C, fn="sum", ngroups=numClasses);

  # Compute the total feature count for each class
  # and add the number of features to this sum
  # for subsequent regularization (Laplace's rule)
  classSums = rowSums(classFeatureCounts) + numFeatures*laplaceCorrection;

  # Compute class conditional probabilities
  classConditionals = (classFeatureCounts + laplaceCorrection) / classSums;

  # Compute class priors
  classCounts = aggregate(target=C, groups=C, fn="count", ngroups=numClasses);
  prior = classCounts / numRows;

  # Compute accuracy on training set
  if( verbose ) {
    logProbs = D %*% t(log(classConditionals)) + t(log(prior));
    acc = sum(rowIndexMax(logProbs) == C) / numRows * 100;
    print("Training Accuracy (%): " + acc);
  }
}

