package com.aliyun.odps.spark.ad.xgboost.v20240808

import com.aliyun.odps.spark.examples.myUtils.ParamUtils
import ml.dmlc.xgboost4j.scala.spark.XGBoostClassifier
import org.apache.commons.lang.StringUtils
import org.apache.commons.lang3.math.NumberUtils
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.types.{DataTypes, StructField}
import org.apache.spark.sql.{Dataset, Row, SparkSession}

import java.net.URL
import java.time.LocalDateTime
import java.time.format.DateTimeFormatter
import scala.io.Source

object XGBoostTrain {
  def main(args: Array[String]): Unit = {
    try {

      val param = ParamUtils.parseArgs(args)

      val dt = LocalDateTime.now.format(DateTimeFormatter.ofPattern("yyyyMMddHHmmSS"))

      val spark = SparkSession.builder()
        .appName("XGBoostTrain:" + dt)
        .getOrCreate()
      val sc = spark.sparkContext

      val loader = getClass.getClassLoader

      val readPath = param.getOrElse("trainReadPath", "")
      val predictReadPath = param.getOrElse("predictReadPath", "")
      val filterNameSet = param.getOrElse("filterNames", "").split(",").filter(_.nonEmpty).toSet
      val featureNameFile = param.getOrElse("featureNameFile", "20240718_ad_feature_name.txt")

      val featureNameContent = readFile(loader.getResource(featureNameFile))

      val featureNameList: List[String] = featureNameContent.split("\n")
        .map(r => r.replace(" ", "").replaceAll("\n", ""))
        .filter(r => r.nonEmpty)
        .filter(r => !containsAny(filterNameSet, r))
        .toList

      val rowRDD = dataMap(sc.textFile(readPath), featureNameList)

      println(s"rowRDD count ${rowRDD.count()}")

      val fields: Array[StructField] = Array(
        DataTypes.createStructField("label", DataTypes.IntegerType, true)
      ) ++ featureNameList.map(f => DataTypes.createStructField(f, DataTypes.DoubleType, true))

      val trainDataSet: Dataset[Row] = spark.createDataFrame(rowRDD, DataTypes.createStructType(fields))

      val vectorAssembler = new VectorAssembler().setInputCols(featureNameList.toArray).setOutputCol("features")

      val xgbInput = vectorAssembler.transform(trainDataSet).select("features", "label")
      xgbInput.show()

      // 创建 XGBoostClassifier 对象
      val xgbClassifier = new XGBoostClassifier()
        .setEta(0.01f)
        .setMissing(0.0f)
        .setMaxDepth(5)
        .setNumRound(1000)
        .setSubsample(0.8)
        .setColsampleBytree(0.8)
        .setScalePosWeight(1)
        .setObjective("binary:logistic")
        .setEvalMetric("auc")
        .setFeaturesCol("features")
        .setLabelCol("label")
        .setNthread(1)
        .setNumWorkers(22)

      // 训练模型
      val model = xgbClassifier.fit(xgbInput)



    }
    catch {
      case e: Throwable => e.printStackTrace()
    }
  }

  private def readFile(filePath: URL): String = {
    var source: Option[Source] = None
    try {
      source = Some(Source.fromURL(filePath))
      return source.get.getLines().mkString("\n")
    }
    catch {
      case e: Exception => println("文件读取异常: " + e.toString)
    }
    finally {
      source.foreach(_.close())
    }
    ""
  }

  private def containsAny(list: Iterable[String], s: String): Boolean = {
    for (item <- list) {
      if (s.contains(item)) {
        return true
      }
    }
    false
  }

  private def dataMap(data: RDD[String], featureNameList: List[String]): RDD[Row] = {
    data.map(r => {
      val line: Array[String] = StringUtils.split(r, "\t")
      val label: Int = NumberUtils.toInt(line(0))

      val map: Map[String, Double] = line.drop(1).map { entry =>
        val Array(key, value) = entry.split(":")
        key -> NumberUtils.toDouble(value, 0.0)
      }.toMap

      val v: Array[Any] = Array.ofDim[Any](featureNameList.length + 1)
      v(0) = label

      for (index <- featureNameList.indices) {
        v(index + 1) = map.getOrElse(featureNameList(index), 0.0)
      }

      Row.fromSeq(v)
    })
  }
}
