package org.genericsystem.reinforcer;

import java.io.IOException;
import java.nio.file.Paths;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.Date;

import org.apache.spark.SparkConf;
import org.apache.spark.ml.Pipeline;
import org.apache.spark.ml.PipelineStage;
import org.apache.spark.ml.classification.LinearSVC;
import org.apache.spark.ml.classification.OneVsRest;
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
import org.apache.spark.ml.feature.HashingTF;
import org.apache.spark.ml.feature.IDF;
import org.apache.spark.ml.feature.IndexToString;
import org.apache.spark.ml.feature.StringIndexer;
import org.apache.spark.ml.feature.Tokenizer;
import org.apache.spark.ml.param.ParamMap;
import org.apache.spark.ml.tuning.CrossValidator;
import org.apache.spark.ml.tuning.CrossValidatorModel;
import org.apache.spark.ml.tuning.ParamGridBuilder;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.api.java.UDF1;
import org.apache.spark.sql.types.DataTypes;

public class SVMClassifier {

	public static void main(String[] args) {
		SparkConf sparkConf = new SparkConf().setMaster("local[2]");
		final SparkSession spark = SparkSession
				.builder()
				.config(sparkConf)
				.appName("SVMClassifier")
				.getOrCreate();
		new SVMClassifier().loadData(spark);
		spark.stop();
	}

	public static class LabelComputer {
		public String computeLabel(String fileName) {
			return Paths.get(fileName).getParent().getFileName().toString();
		}
	}

	public void loadData(SparkSession spark) {
		UDF1<String, String> labelUDF = t1 -> new LabelComputer().computeLabel(t1);
		spark.udf().register("label", labelUDF, DataTypes.StringType);
		Dataset<Row> data = spark.read().text("pieces/text/*").withColumn("label", org.apache.spark.sql.functions.input_file_name());
		data = data.withColumn("label", org.apache.spark.sql.functions.callUDF("label", data.col("label")));

		Dataset<Row>[] splits = data.randomSplit(new double[] {0.7, 0.3});
		Dataset<Row> trainData = splits[0];
		Dataset<Row> validData = splits[1];

		StringIndexer indexer = new StringIndexer()
				.setInputCol("label")
				.setOutputCol("labelIndex")
				.setHandleInvalid("keep");
		Tokenizer tokenizer = new Tokenizer()
				.setInputCol("value")
				.setOutputCol("words");
		HashingTF hashingTF = new HashingTF()
				.setInputCol(tokenizer.getOutputCol())
				.setOutputCol("rawFeatures")
				.setNumFeatures(20);
		IDF idf = new IDF()
				.setInputCol(hashingTF.getOutputCol())
				.setOutputCol("features");
		LinearSVC lsvc = new LinearSVC()
				.setMaxIter(100)
				.setRegParam(0.1);
		OneVsRest ovr = new OneVsRest()
				.setClassifier(lsvc)
				.setFeaturesCol(idf.getOutputCol())
				.setLabelCol(indexer.getOutputCol());
		IndexToString converter = new IndexToString()
				.setInputCol("prediction")
				.setOutputCol("origPrediction");

		Pipeline pipeline = new Pipeline()
				.setStages(new PipelineStage[] { indexer, tokenizer, hashingTF, idf, ovr, converter });

		ParamMap[] paramGrid = new ParamGridBuilder()
				.addGrid(hashingTF.numFeatures(), new int[] { 20 , 50, 200, 1000 })
				.addGrid(lsvc.maxIter(), new int[] { 20, 50, 100, 200 })
				.addGrid(lsvc.regParam(), new double[] { 0.05, 0.1 })
				.build();

		MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator()
				.setMetricName("accuracy")
				.setLabelCol(indexer.getOutputCol());

		CrossValidator cv = new CrossValidator()
				.setEstimator(pipeline)
				.setEvaluator(evaluator)
				.setEstimatorParamMaps(paramGrid)
				.setNumFolds(3);

		CrossValidatorModel cvModel = cv.fit(trainData);

		Dataset<Row> predictions = cvModel.transform(validData);
		double accuracy = evaluator.evaluate(predictions);
		System.out.println("Test error = " + (1 - accuracy));

		DateFormat format = new SimpleDateFormat("yyyyMMddHHmmss");
		try {
			cvModel.save("SVMModel-" + format.format(new Date()));
		} catch (IOException e) {
			throw new RuntimeException("Exception while saving trained model", e);
		}
	}	
}

@@ -43,28 +43,124 @@
 ort 
-java.io.IOException;
+static org.apache.spark.sql.functions.input_file_name;%0Aimport static org.apache.spark.sql.functions.regexp_extract;%0A
 %0Aimp
@@ -172,22 +172,22 @@
 ava.
-n
 io.
-file.Paths
+IOException
 ;%0Aim
@@ -1112,96 +1112,8 @@
 ion;
-%0Aimport org.apache.spark.sql.api.java.UDF1;%0Aimport org.apache.spark.sql.types.DataTypes;
 %0A%0Apu
@@ -1446,347 +1446,64 @@
 lic 
-static class LabelComputer %7B%0A%09%09public String computeLabel(String fileName) %7B%0A%09%09%09return Paths.get(fileName).getParent().getFileName().toString();%0A%09%09%7D%0A%09%7D%0A%0A%09public void loadData(SparkSession spark) %7B%0A%09%09UDF1%3CString, String%3E labelUDF = t1 -%3E new LabelComputer().computeLabel(t1);%0A%09%09spark.udf().register(%22label%22, labelUDF, DataTypes.StringType);
+Dataset%3CRow%3E loadData(SparkSession spark, String path) %7B
 %0A%09%09D
@@ -1543,23 +1543,12 @@
 ext(
-%22pieces/text/*%22
+path
 ).wi
@@ -1569,39 +1569,8 @@
 l%22, 
-org.apache.spark.sql.functions.
 inpu
@@ -1587,22 +1587,22 @@
 e());%0A%09%09
-data =
+return
  data.wi
@@ -1623,77 +1623,180 @@
 l%22, 
-org.apache.spark.sql.functions.callUDF(%22label%22, data.col(%22label%22)));%0A
+regexp_extract(data.col(%22label%22), %22.*/(%5B%5E/%5D*)/%5B%5E/%5D*%22, 1)).cache();%0A%09%7D%0A%0A%09public void trainModel(SparkSession spark) %7B%0A%09%09Dataset%3CRow%3E data = loadData(spark, %22pieces/text/*%22);
 %0A%09%09D

Remove UDF, use sql.functions instead