#' @include FamiliarS4Generics.R
#' @include FamiliarS4Classes.R
NULL



# familiarNaiveBayes -----------------------------------------------------------
setClass(
  "familiarNaiveBayes",
  contains = "familiarModel"
)



# initialize -------------------------------------------------------------------
setMethod(
  "initialize",
  signature(.Object = "familiarNaiveBayes"),
  function(.Object, ...) {
    # Update with parent class first.
    .Object <- callNextMethod()

    # Set required package
    .Object@package <- "e1071"

    return(.Object)
  }
)



# is_available -----------------------------------------------------------------
setMethod(
  "is_available",
  signature(object = "familiarNaiveBayes"),
  function(object, ...) {
    # Naive bayes is only available for categorical outcomes.
    if (object@outcome_type %in% c("binomial", "multinomial")) {
      return(TRUE)
    } else {
      return(FALSE)
    }
  }
)



# get_default_hyperparameters --------------------------------------------------
setMethod(
  "get_default_hyperparameters",
  signature(object = "familiarNaiveBayes"),
  function(object, data = NULL, ...) {
    # Initialise list and declare hyperparameter entries.
    param <- list()
    param$sign_size <- list()
    param$laplace <- list()

    # If data is explicitly set to NULL, return the list with
    # hyperparameter names only.
    if (is.null(data)) return(param)
    
    # signature size -----------------------------------------------------------
    param$sign_size <- .get_default_sign_size(
      data = data, 
      restrict_samples = TRUE
    )

    # laplace smoothing parameter ----------------------------------------------
    param$laplace <- .set_hyperparameter(
      default = c(0.0, 1.0, 2.0, 5.0),
      type = "numeric",
      range = c(0.0, 10.0),
      valid_range = c(0.0, Inf),
      randomise = TRUE
    )

    return(param)
  }
)



# ..train ----------------------------------------------------------------------
setMethod(
  "..train",
  signature(
    object = "familiarNaiveBayes",
    data = "dataObject"
  ),
  function(object, data, ...) {
    # Check if training data is ok.
    if (reason <- has_bad_training_data(object = object, data = data)) {
      return(callNextMethod(object = .why_bad_training_data(
        object = object,
        reason = reason
      )))
    }

    # Check if hyperparameters are set.
    if (is.null(object@hyperparameters)) {
      return(callNextMethod(object = ..update_errors(
        object = object,
        ..error_message_no_optimised_hyperparameters_available()
      )))
    }

    # Check that required packages are loaded and installed.
    require_package(object, "train")

    # Find feature columns in the data.
    feature_columns <- get_feature_columns(x = data)

    # Parse formula
    formula <- stats::reformulate(
      termlabels = feature_columns,
      response = quote(outcome)
    )

    # Train the model.
    model <- do.call_with_handlers(
      e1071::naiveBayes,
      args = list(
        formula,
        "data" = data@data,
        "laplace" = object@hyperparameters$laplace
      )
    )

    # Extract values.
    object <- ..update_warnings(object = object, model$warning)
    object <- ..update_errors(object = object, model$error)
    model <- model$value

    # Check if the model trained at all.
    if (!is.null(object@messages$error)) {
      return(callNextMethod(object = object))
    }

    # Add model to the familiarModel object.
    object@model <- model

    # Set learner version
    object <- set_package_version(object)

    return(object)
  }
)



# ..train_naive ----------------------------------------------------------------
setMethod(
  "..train_naive",
  signature(
    object = "familiarNaiveBayes",
    data = "dataObject"
  ),
  function(object, data, ...) {
    if (object@outcome_type %in% c("binomial", "multinomial")) {
      # Turn into a Naive model.
      object <- methods::new("familiarNaiveModel", object)
    }

    return(..train(
      object = object,
      data = data,
      ...
    ))
  }
)



# ..predict --------------------------------------------------------------------
setMethod(
  "..predict",
  signature(
    object = "familiarNaiveBayes",
    data = "dataObject"
  ),
  function(object, data, type = "default", ...) {
    # Check that required packages are loaded and installed.
    require_package(object, "predict")

    # Check if the model was trained.
    if (!model_is_trained(object)) {
      return(callNextMethod())
    }
    
    # Check if the data is empty.
    if (is_empty(data)) {
      return(callNextMethod())
    }
    
    if (type == "default") {
      # default ----------------------------------------------------------------
      
      if (object@outcome_type %in% c("binomial", "multinomial")) {
        
        # Use the model to predict class probabilities.
        model_predictions <- predict(
          object = object@model,
          newdata = data@data,
          type = "raw"
        )
        
        # Obtain class levels.
        class_levels <- get_outcome_class_levels(x = object)
        
        # Add class probabilities.
        prediction_list <- list()
        for (ii in seq_along(class_levels)) {
          prediction_list[[class_levels[ii]]] <- model_predictions[, class_levels[ii]]
        }
        
        # Store as prediction table.
        prediction_table <- as_prediction_table(
          x = prediction_list,
          type = "classification",
          data = data,
          model_object = object
        )
        
      } else {
        ..error_outcome_type_not_implemented(object@outcome_type)
      }

      return(prediction_table)
      
    } else if (!.is_available_prediction_type(type)) {
      # user-specified method --------------------------------------------------

      # Use the model for prediction.
      return(predict(
        object = object@model,
        newdata = data@data,
        type = type,
        ...
      ))
      
    } else {
      ..error_no_predictions_possible(object, type)
    }
  }
)



# ..vimp------------------------------------------------------------------------
# Naive Bayes does not have an associated variable importance method.



# .trim_model-------------------------------------------------------------------
setMethod(
  ".trim_model",
  signature(object = "familiarNaiveBayes"),
  function(object, ...) {
    # Update model by removing the call.
    object@model$call <- call("trimmed")

    # Add show.
    object <- .capture_show(object)

    # Set is_trimmed to TRUE.
    object@is_trimmed <- TRUE

    # Default method for models that lack a more specific method.
    return(object)
  }
)



.get_available_naive_bayes_learners <- function(show_general = TRUE) {
  return("naive_bayes")
}
