package org.jeecg.cyl

import org.apache.spark.sql.SparkSession
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.regression.LinearRegression
import org.apache.spark.ml.Pipeline
import org.apache.spark.sql.catalyst.dsl.expressions.StringToAttributeConversionHelper
object forecast {
  def main(args: Array[String]): Unit = {
    import org.apache.spark.sql.SparkSession
    import org.apache.spark.sql.functions._

    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\\cyl\\tianchi_mum_baby_trade_history.csv")

    val filteredData = data.filter(col("category_1").equalTo(50013636))

    // 准备特征和标签
    val assembler = new VectorAssembler()
      .setInputCols(Array("day"))  // 可以添加更多的特征列
      .setOutputCol("features")

    // 定义线性回归模型
    val lr = new LinearRegression()
      .setLabelCol("buy_mount")  // 替换为你的销量列名
      .setFeaturesCol("features")

    // 创建pipeline
    val pipeline = new Pipeline().setStages(Array(assembler, lr))

    // 训练模型
    val model = pipeline.fit(filteredData)

    // 进行预测
    val predictions = model.transform(filteredData)

    // 显示预测结果
    predictions.select("day", "buy_mount", "prediction").show()

    spark.stop()
  }
}
