using ScikitLearn
using Missings
#using Meta
#using Mocha
@sk_import linear_model: LogisticRegression
@sk_import svm: LinearSVC
@sk_import svm: SVC
@sk_import svm: NuSVC
@sk_import ensemble: RandomForestClassifier
@sk_import neural_network: MLPClassifier
include("crossval_and_stats.jl")
#include("mocha_helper.jl")


#####################################################################
#                         CLASSIFIER
#####################################################################

function model_and_predict_skl(train_features, train_labels, test_features, param_dict, classifier)

  #parameter for prediction
  proba = get_param(param_dict, :proba, false)

  #fit model on training data
  println("$(string(classifier)) (SKL) - Build Model - Begin")
  model = skl_model(param_dict, classifier)
  ScikitLearn.fit!(model, train_features, train_labels)
  println("$(string(classifier)) (SKL) - Build Model - Complete")

  #make predictions on testing data
  println("$(string(classifier)) (SKL) - Apply Model - Begin")
  predictions = skl_predict(model, test_features, proba)
  println("$(string(classifier)) (SKL) - Apply Model - Complete")

  return predictions

end

#####################################################################
#                        HELPER FUNCTIONS
#####################################################################

function skl_model(param_dict, classifier)

  if classifier == :LogisticRegression
    penalty = get_param(param_dict, :penalty, "l2")
    dual = get_param(param_dict, :dual, false)
    C = get_param(param_dict, :C, 1.0)
    fit_intercept = get_param(param_dict, :fit_intercept, true)
    intercept_scaling = get_param(param_dict, :intercept_scaling, 1)
    class_weight = get_param(param_dict, :class_weight, nothing)
    max_iter = get_param(param_dict, :max_iter, 100)
    random_state = get_param(param_dict, :random_state, nothing)
    solver = get_param(param_dict, :solver, "liblinear")
    tol = get_param(param_dict, :tol, 0.0001)
    verbose = get_param(param_dict, :verbose, 0)
    warm_start = get_param(param_dict, :warm_start, false)
    n_jobs = get_param(param_dict, :n_jobs, 1)
    #not included: multi_class

    return LogisticRegression(penalty=penalty, dual=dual, C=C, fit_intercept=fit_intercept, intercept_scaling=intercept_scaling, class_weight=class_weight, max_iter=max_iter, random_state=random_state, solver=solver, tol=tol, verbose=verbose, warm_start=warm_start, n_jobs=n_jobs)

  elseif classifier == :LinearSVC
    C = get_param(param_dict, :C, 1.0)
    loss = get_param(param_dict, :loss, "squared_hinge")
    penalty = get_param(param_dict, :penalty, "l2")
    dual = get_param(param_dict, :dual, true)
    tol = get_param(param_dict, :tol, 0.0001)
    fit_intercept = get_param(param_dict, :fit_intercept, true)
    intercept_scaling = get_param(param_dict, :intercept_scaling, 1)
    class_weight = get_param(param_dict, :class_weight, nothing)
    verbose = get_param(param_dict, :verbose, 0)
    random_state = get_param(param_dict, :random_state, nothing)
    max_iter = get_param(param_dict, :max_iter, 1000)
    #not included: multi_class

    return LinearSVC(C=C, loss=loss, penalty=penalty, dual=dual, tol=tol, fit_intercept=fit_intercept, intercept_scaling=intercept_scaling, class_weight=class_weight, verbose=verbose, random_state=random_state, max_iter=max_iter)

  elseif classifier == :SVC
    C = get_param(param_dict, :C, 1.0)
    kernel = get_param(param_dict, :kernal, "rbf")
    degree = get_param(param_dict, :degree, 3)
    gamma = get_param(param_dict, :gamma, "auto")
    coef0 = get_param(param_dict, :coef0, 0.0)
    probability = get_param(param_dict, :proba, false)
    shrinking = get_param(param_dict, :shrinking, true)
    tol = get_param(param_dict, :tol, 0.001)
    class_weight = get_param(param_dict, :class_weight, nothing)
    verbose = get_param(param_dict, :verbose, false)
    max_iter = get_param(param_dict, :max_iter, -1)
    decision_function_shape = get_param(param_dict, :decision_function_shape, nothing)
    random_state = get_param(param_dict, :random_state, nothing)
    #not included: cache_size

    return SVC(C=C, kernel=kernel, degree=degree, gamma=gamma, coef0=coef0, probability=probability, shrinking=shrinking, tol=tol, class_weight=class_weight, verbose=verbose, max_iter=max_iter, decision_function_shape=decision_function_shape, random_state=random_state)

  elseif classifier == :NuSVC
    nu = get_param(param_dict, :nu, 0.5)
    kernel = get_param(param_dict, :kernal, "rbf")
    degree = get_param(param_dict, :degree, 3)
    gamma = get_param(param_dict, :gamma, "auto")
    coef0 = get_param(param_dict, :coef0, 0.0)
    probability = get_param(param_dict, :probability, false)
    shrinking = get_param(param_dict, :shrinking, true)
    tol = get_param(param_dict, :tol, 0.001)
    class_weight = get_param(param_dict, :class_weight, nothing)
    verbose = get_param(param_dict, :verbose, false)
    max_iter = get_param(param_dict, :max_iter, -1)
    decision_function_shape = get_param(param_dict, :decision_function_shape, nothing)
    random_state = get_param(param_dict, :random_state, nothing)
    #not included: cache_size

    return NuSVC(nu=nu, kernel=kernel, degree=degree, gamma=gamma, coef0=coef0, probability=probability, shrinking=shrinking, tol=tol, class_weight=class_weight, min_impurity_decrease=min_impurity_decrease, verbose=verbose, max_iter=max_iter, decision_function_shape=decision_function_shape, random_state=random_state)

  elseif classifier == :RandomForest
    n_estimators = get_param(param_dict, :n_estimators, 10)
    criterion = get_param(param_dict, :criterion, "gini")
    max_features = get_param(param_dict, :max_features, "auto")
    max_depth = get_param(param_dict, :max_depth, nothing)
    min_samples_split = get_param(param_dict, :min_samples_split, 6)
    min_samples_leaf = get_param(param_dict, :min_samples_leaf, 1)
    min_weight_fraction_leaf = get_param(param_dict, :min_weight_fraction_leaf, 0.)
    max_leaf_nodes = get_param(param_dict, :max_leaf_nodes, nothing)
    min_impurity_decrease = get_param(param_dict, :min_impurity_decrease, 0.00000001)
    bootstrap = get_param(param_dict, :bootstrap, true)
    oob_score = get_param(param_dict, :oob_score, false)
    n_jobs = get_param(param_dict, :n_jobs, 1)
    random_state = get_param(param_dict, :random_state, nothing)
    verbose = get_param(param_dict, :verbose, 0)
    warm_start = get_param(param_dict, :warm_start, false)
    class_weight = get_param(param_dict, :class_weight, nothing)

    return RandomForestClassifier(n_estimators=n_estimators, criterion=criterion, max_features=max_features, max_depth=max_depth, min_samples_split=min_samples_split, min_samples_leaf=min_samples_leaf, min_weight_fraction_leaf=min_weight_fraction_leaf, max_leaf_nodes=max_leaf_nodes, bootstrap=bootstrap, oob_score=oob_score, n_jobs=n_jobs, random_state=random_state, verbose=verbose, warm_start=warm_start, class_weight=class_weight)

  elseif classifier == :MLP
    hidden_layer_sizes = get_param(param_dict, :hidden_layer_sizes, (100,))
    activation = get_param(param_dict, :activation, "relu")
    solver = get_param(param_dict, :solver, "adam")
    alpha = get_param(param_dict, :alpha, 0.0001)
    batch_size = get_param(param_dict, :batch_size, "auto")
    learning_rate = get_param(param_dict, :learning_rate, "constant")
    max_iter = get_param(param_dict, :max_iter, 200)
    random_state = get_param(param_dict, :random_state, nothing)
    shuffle = get_param(param_dict, :shuffle, true)
    tol = get_param(param_dict, :tol, 0.0001)
    learning_rate_init = get_param(param_dict, :learning_rate_init, 0.001)
    power_t = get_param(param_dict, :power_t, 0.5)
    verbose = get_param(param_dict, :verbose, false)
    warm_start = get_param(param_dict, :warm_start, false)
    momentum = get_param(param_dict, :momentum, 0.9)
    nesterovs_momentum = get_param(param_dict, :nesterovs_momentum, true)
    early_stopping = get_param(param_dict, :early_stopping, false)
    validation_fraction = get_param(param_dict, :validation_fraction, 0.1)
    beta_1 = get_param(param_dict, :beta_1, 0.9)
    beta_2 = get_param(param_dict, :beta_2, 0.999)
    epsilon = get_param(param_dict, :epsilon, 0.00000001)
    #not included: min_impurity_split, n_jobs, warm_start

    return MLPClassifier(hidden_layer_sizes=hidden_layer_sizes, activation=activation, solver=solver, alpha=alpha, batch_size=batch_size, learning_rate=learning_rate, max_iter=max_iter, random_state=random_state, shuffle=shuffle, tol=tol, learning_rate_init=learning_rate_init, power_t=power_t, verbose=verbose, warm_start=warm_start, momentum=momentum, nesterovs_momentum=nesterovs_momentum, early_stopping=early_stopping, validation_fraction=validation_fraction, beta_1=beta_1, beta_2=beta_2, epsilon=epsilon)

  else
    error("classifier $(string(classifier)) is not supported")
  end
end


function get_param(dict, key, default)
  if haskey(dict, key)
    return dict[key]
  else
    return default
  end
end


#calls skl predictoin function according to proba paramter
function skl_predict(model, test_features, proba=false)
  predictions = []
  if proba == false
    return ScikitLearn.predict(model, test_features)
  else
    return ScikitLearn.predict_proba(model, test_features)[:,2]
  end
end


function get_features_labels(df, pmids)

  i = 0

  all_pmids = df[:pmid]
  indices = falses(length(all_pmids))
  for i in 1:length(all_pmids)
    if all_pmids[i] in pmids
      indices[i] = true
    end
  end

  temp_pmids = all_pmids[indices]
  labels = df[indices,:class]
  features = convert(Array,df[indices, 4:end])
  temp_genes = df[indices,:accepted_genes]

  genes = fill([],length(temp_genes))

  for i in 1:length(temp_genes)
    g = temp_genes[i]
    if !ismissing(g)
      genes[i] = Meta.parse.(split(g,"|"))
    end
  end

  return features, labels, genes
end


#####################################################################
#                       NON-SKL CLASSIFIERS
#####################################################################

#=
function MLPmocha(train_features, train_labels, test_features, param_dict)

  proba = get_param(param_dict, :proba, false)

  backend = CPUBackend()
  init(backend)
  train_features = convert(Array{Float64,2}, train_features)
  test_features = convert(Array{Float64,2}, test_features)
  train_labels = convert(Array{Float64,1}, train_labels)
  train_features = float(convert_rows_to_columns(train_features))
  train_labels = float(convert_rows_to_columns(train_labels))
  test_features = float(convert_rows_to_columns(test_features))

  train_net, test_net = create_nets(backend, train_features, train_labels, test_features, param_dict)

  println("MLP (Mocha) - Build Model - Begin")
  solver = create_solver(param_dict)
  model = solve(solver, train_net)
  println("MLP (Mocha) - Build Model - Complete")

  println("MLP (Mocha) - Apply Model - Begin")
  forward(test_net)
  predictions = get_mocha_predictions(test_net, proba)
  println("MLP (Mocha) - Apply Model - Complete")

  destroy(train_net)
  destroy(test_net)
  shutdown(backend)

  return predictions
end

function get_mocha_train_features_labels(df, pmids)
  features, labels = get_features_labels(df, pmids)
  n_features = float(convert_rows_to_columns(features))
  n_labels = float(convert_rows_to_columns(labels))
  return n_features, n_labels
end

function get_mocha_test_features_labels(df, pmids)
  features, labels = get_features_labels(df, pmids)
  n_features = float(convert_rows_to_columns(features))
  return n_features, labels
end
=#