package org.jeecg.yw.spark

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions.udf
import org.apache.spark.ml.feature.{VectorAssembler, StandardScaler}
import org.apache.spark.ml.classification.LogisticRegression
import org.apache.spark.ml.{Pipeline, PipelineModel}
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator

object train {
  def main(args: Array[String]): Unit = {
    // 设置HADOOP_USER_NAME环境变量
    System.setProperty("HADOOP_USER_NAME", "root")
    // 创建Spark会话(win本地)
    val spark = SparkSession.builder()
      .appName("Weather Temperature Prediction")
      .master("local[*]")
      .getOrCreate()
    // 加载数据
    val data = spark.read.option("header", "true").option("inferSchema", "true")
      .csv("C:\\Users\\23972\\Desktop\\Completion-template-cmd\\system\\start\\src\\main\\java\\org\\jeecg\\yw\\spark\\青少年膳食营养数据集.csv")
    // 将健康评分转为分类标签的UDF
    val categorizeHealth = udf((score: Double) => score match {
      case score if score <= 33 => 0
      case score if score <= 66 => 1
      case _ => 2
    })
    // 应用UDF
    val processedData = data.withColumn("HealthCategory", categorizeHealth(data("健康自评")))
    val assembler = new VectorAssembler()
      .setInputCols(Array("绿叶蔬菜", "水果", "大豆制品", "新鲜肉类"))
      .setOutputCol("features")
    // 特征标准化
    val scaler = new StandardScaler()
      .setInputCol("features")
      .setOutputCol("scaledFeatures")
    // 初始化逻辑回归模型
    val lr = new LogisticRegression()
      .setFeaturesCol("scaledFeatures")
      .setLabelCol("HealthCategory")
    // 创建Pipeline
    val pipeline = new Pipeline()
      .setStages(Array(assembler, scaler, lr))
    // 划分数据集
    val Array(trainingData, testData) = processedData.randomSplit(Array(0.8, 0.2), seed = 1234L)
    // 训练模型
    val model = pipeline.fit(trainingData)
    // 预测
    val predictions = model.transform(testData)
    // 模型评估
    val evaluator = new MulticlassClassificationEvaluator()
      .setLabelCol("HealthCategory")
      .setPredictionCol("prediction")
      .setMetricName("accuracy")

    val accuracy = evaluator.evaluate(predictions)
    println(s"Accuracy = $accuracy")
    // 保存在 hdfs路径
    model.write.overwrite().save("hdfs://192.168.192.100:8020/Model")
    // 关闭Spark会话
    spark.stop()
  }
}
