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

# Implements multinomial naive Bayes classifier with Laplace correction
#
# Example Usage:
# hadoop jar SystemDS.jar -f naive-bayes.dml -nvargs X=<Data> Y=<labels> laplace=<Laplace Correction> prior=<Model file1> conditionals=<Model file2> accuracy=<accuracy file> fmt="text"
#

# defaults
laplaceCorrection = ifdef($laplace, 1)
fmt = ifdef($fmt, "text")

# reading input args
D = read($X)
C = read($Y)
numRows = nrow(D)
numFeatures = ncol(D)
minFeatureVal = min(D)
numClasses = max(C)
minLabelVal = min(C)

# sanity checks of data and arguments
if(minFeatureVal < 0)
	stop("Stopping due to invalid argument: Multinomial naive Bayes is meant for count-based feature values, minimum value in X is negative")
if(numRows < 2)
	stop("Stopping due to invalid inputs: Not possible to learn a classifier without at least 2 rows")
if(minLabelVal < 1)
	stop("Stopping due to invalid argument: Label vector (Y) must be recoded")
if(numClasses == 1)
	stop("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("Stopping due to invalid argument: Please ensure that Y contains (positive) integral labels")
if(laplaceCorrection < 0)
	stop("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=as.integer(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=as.integer(numClasses))
classPrior = classCounts / numRows;

# Compute accuracy on training set
logProbs = D %*% t(log(classConditionals)) + t(log(classPrior));
acc = sum(rowIndexMax(logProbs) == C) / numRows * 100

acc_str = "Training Accuracy (%): " + acc
print(acc_str)
accuracyFile = $accuracy
if(accuracyFile != " ") {
	write(acc, accuracyFile)
}

extraModelParams = as.matrix(numFeatures)
classPrior = rbind(classPrior, extraModelParams)

# write out the model
write(classPrior, $prior, format=fmt);
write(classConditionals, $conditionals, format=fmt);