package com.example.model.classification
import java.util
import com.example.SparkConnect
import com.example.entity.Model
import com.example.model.helper.Utils
import org.apache.spark.ml.classification.{DecisionTreeClassificationModel, DecisionTreeClassifier, LogisticRegression, LogisticRegressionModel}
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.ml.feature.{IndexToString, StringIndexer, VectorIndexer}
import org.apache.spark.ml.{Pipeline, PipelineModel, PipelineStage}
import org.apache.spark.sql.{Dataset, Row}
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.stereotype.Component
import com.example.util.Constant

import scala.collection.mutable

/**
 * 分类方法的默认实现
 * @author   fanyanpeng
 * @date 2023/3/22 16:15
 */
class DefaultClassification extends SparkConnect {


  @Autowired
  val utils: Utils = null

  @throws(classOf[Exception])
  def training(libsvmFileUri: String,
               params: util.Map[String,Object]): (PipelineModel,util.Map[String, String]) = {

    val trainSet = sparkSession.read.format("libsvm").load(libsvmFileUri)
    trainSet.show()

    val maxCategories = params.get(Constant.TRAIN_PARAM_MAX_CATEGORY).asInstanceOf[Int]
    val featureIndexer = new VectorIndexer()
      .setInputCol("features")
      .setOutputCol("indexedFeatures")
      .setMaxCategories(maxCategories)
      .fit(trainSet)

    trainSet.show()

    // Split the data into training and test sets (30% held out for testing).
    //为了对训练效果进行评估，将训练集进一步划分为子训练集与子测试集
    var Array(subTrainSet:Dataset[Row], subTestSet:Dataset[Row]) = trainSet.randomSplit(Array(0.7, 1 - 0.7))

    /*******************************************/
    //需要由子类实现
    val classifier = buildClassifier(params)
    /*******************************************/

    val pipeline = new Pipeline()
      .setStages(Array(featureIndexer, classifier))

    // Train model. This also runs the indexers.
    val pipelineModel:PipelineModel = pipeline.fit(subTrainSet)


    val resultMap:util.Map[String,String] = new util.HashMap[String,String]()
    // Make predictions.
    val predictions = pipelineModel.transform(subTestSet)

    // Select (prediction, true label) and compute test error.
    //为了简便，使用了多分类
    val evaluator = new MulticlassClassificationEvaluator()
      .setLabelCol("label")
      .setPredictionCol("prediction")


    val weightedRecall = evaluator.setMetricName("weightedRecall").evaluate(predictions);
    val accuracy = evaluator.setMetricName("accuracy").evaluate(predictions)
    log.info("Test Error = " + (1.0 - accuracy))
    val weightedPrecision = evaluator.setMetricName("weightedPrecision").evaluate(predictions);
    val f1 = evaluator.setMetricName("f1").evaluate(predictions);


    val model = pipelineModel.stages(1).asInstanceOf[LogisticRegressionModel]
    //println("Learned classification tree model:\n" + treeModel.toDebugString)
    log.info("\n\nLearned classification Logistic model:\n" + model.toString());


    resultMap.put("Accuracy: ", accuracy.toString)
    resultMap.put("weightedPrecision", weightedPrecision.toString)
    resultMap.put("weightedRecall", weightedRecall.toString)
    resultMap.put("f1", f1.toString)

    return (pipelineModel, resultMap)
  }


  /**
   * 默认实现：逻辑回归
   * @author   fanyanpeng
   * @date 2023/3/22 16:44
   * @param params 参数列表
   * @return org.apache.spark.ml.PipelineStage
   */
  def buildClassifier(params:util.Map[String,Object]):PipelineStage  = {
    val maxIter = params.get(Constant.TRAIN_PARAM_MAX_ITER).asInstanceOf[Int]
    val regParam = params.get(Constant.TRAIN_PARAM_REG_PARAM).asInstanceOf[Float]
    val elasticNetParam = params.get(Constant.TRAIN_PARAM_ELASTICNET_PARAM).asInstanceOf[Float]

    val lr = new LogisticRegression()
      .setLabelCol("label")
      .setFeaturesCol("indexedFeatures")
      .setMaxIter(maxIter)
      .setRegParam(regParam)
      .setElasticNetParam(elasticNetParam)
    lr
  }


}