# $Id$
#
#  Copyright (C) 2004-2005 Rational Discovery LLC.
#   All Rights Reserved
#
""" handles doing cross validation with naive bayes models
and evaluation of individual models

"""
from __future__ import print_function
from rdkit.ML.NaiveBayes.ClassificationModel import NaiveBayesClassifier
from rdkit.ML.Data import SplitData
try:
  from rdkit.ML.FeatureSelect import CMIM
except ImportError:
  CMIM=None

def makeNBClassificationModel(trainExamples, attrs, nPossibleValues, nQuantBounds,
                              mEstimateVal=-1.0,
                              useSigs=False,
                              ensemble=None,useCMIM=0,
                              **kwargs) :
  if CMIM is not None and useCMIM > 0 and useSigs and not ensemble:
    ensemble = CMIM.SelectFeatures(trainExamples,useCMIM,bvCol=1)
  if ensemble:
    attrs = ensemble
  model = NaiveBayesClassifier(attrs, nPossibleValues, nQuantBounds, 
                               mEstimateVal=mEstimateVal,useSigs=useSigs)


  model.SetTrainingExamples(trainExamples)
  model.trainModel()
  return model
    
def CrossValidate(NBmodel, testExamples, appendExamples=0) :
    
  nTest = len(testExamples)
  assert nTest,'no test examples: %s'%str(testExamples)
  badExamples = []
  nBad = 0
  preds = NBmodel.ClassifyExamples(testExamples, appendExamples)
  assert len(preds) == nTest

  for i in range(nTest):
    testEg = testExamples[i]
    trueRes = testEg[-1]
    res = preds[i]

    if (trueRes != res) :
      badExamples.append(testEg)
      nBad += 1
  return float(nBad)/nTest, badExamples

def CrossValidationDriver(examples, attrs, nPossibleValues, nQuantBounds,
                          mEstimateVal=0.0,
                          holdOutFrac=0.3, modelBuilder=makeNBClassificationModel,
                          silent=0, calcTotalError=0, **kwargs) :
  nTot = len(examples)
  if not kwargs.get('replacementSelection',0):
    testIndices,trainIndices = SplitData.SplitIndices(nTot,holdOutFrac,
                                                      silent=1,legacy=1,
                                                      replacement=0)
  else :
    testIndices,trainIndices = SplitData.SplitIndices(nTot,holdOutFrac,
                                                      silent=1,legacy=0,
                                                      replacement=1)

  trainExamples = [examples[x] for x in trainIndices]
  testExamples = [examples[x] for x in testIndices]

  NBmodel = modelBuilder(trainExamples, attrs, nPossibleValues, nQuantBounds,
                         mEstimateVal,**kwargs)

  if not calcTotalError:                # 
    xValError, badExamples = CrossValidate(NBmodel, testExamples,appendExamples=1)
  else:
    xValError,badExamples = CrossValidate(NBmodel, examples,appendExamples=0)

  if not silent:
    print('Validation error was %%%4.2f'%(100*xValError))
  NBmodel._trainIndices = trainIndices
  return NBmodel, xValError

