package com.criteo.lookalike.flow.Validation

import org.apache.spark.ml.{Pipeline}
import org.apache.spark.ml.evaluation.Evaluator
import org.apache.spark.ml.param.ParamMap
import org.apache.spark.ml.tuning.{CrossValidator, CrossValidatorModel}
import org.apache.spark.sql.DataFrame

trait CrossValidatorTrait extends Validators[CrossValidatorModel] {

  def evaluator: Evaluator

  def estimator: Pipeline

  def paramGrid: Array[ParamMap]

  override def fit(estimator: Pipeline,
                   evaluator: Evaluator,
                   paramGrid: Array[ParamMap],
                   dfIn: DataFrame): CrossValidatorModel = {

    val validator = new CrossValidator()
      .setEstimator(estimator)
      .setEvaluator(evaluator)
      .setEstimatorParamMaps(paramGrid)
      .setNumFolds(3)

    validator.fit(dfIn)
  }

  override def transform(model: CrossValidatorModel, df: DataFrame): DataFrame = {
    model.transform(df);
  }
}