package org.jeecg.sy.java;

import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.ml.classification.LogisticRegression;
import org.apache.spark.ml.feature.VectorAssembler;
import org.apache.spark.ml.Pipeline;
import org.apache.spark.ml.PipelineModel;
import org.apache.spark.ml.PipelineStage;

public class ShotPrediction {

    public static void main(String[] args) {
        SparkSession spark = SparkSession.builder()
                .appName("Music Shot Prediction")
                .master("local[2]")
                .getOrCreate();

        // 加载数据
        Dataset<Row> data = spark.read()
                .format("csv")
                .option("header", "true")
                .option("inferSchema", "true")
                .load("path_to_your_csv_file");

        // 数据预处理
        data = data.withColumn("label", data.col("result").cast("integer"))
                .withColumn("shot_type_num", data.col("shot_type").cast("integer"))
                .withColumn("distance_ft_num", data.col("distance_ft").cast("integer"))
                .withColumn("lead_num", data.col("lead").cast("integer"));

        // 特征向量化
        VectorAssembler assembler = new VectorAssembler()
                .setInputCols(new String[]{"shot_type_num", "distance_ft_num", "lead_num"})
                .setOutputCol("features");

        // 创建逻辑回归模型
        LogisticRegression lr = new LogisticRegression();

        // 创建管道
        Pipeline pipeline = new Pipeline()
                .setStages(new PipelineStage[]{assembler, lr});

        // 划分数据集
        Dataset<Row>[] splits = data.randomSplit(new double[]{0.7, 0.3}, 42);
        Dataset<Row> trainingData = splits[0];
        Dataset<Row> testData = splits[1];

        // 训练模型
        PipelineModel model = pipeline.fit(trainingData);

        // 模型评估
        Dataset<Row> predictions = model.transform(testData);
        predictions.select("features", "label", "prediction").show();

        // 关闭Spark会话
        spark.stop();
    }
}


