from pyspark.sql import SparkSession
from pyspark.sql.types import DoubleType, IntegerType
from pyspark.sql.functions import col
from pyspark.ml.feature import VectorAssembler
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.evaluation import MulticlassClassificationEvaluator

spark = SparkSession.builder \
    .appName("IDS-Distributed-Training") \
    .master("spark://172.29.5.41:7077") \
    .getOrCreate()

spark.sparkContext.setLogLevel("WARN")
print("Spark Session started. Master: " + spark.sparkContext.master)

# Read data
df = spark.read.csv(
    "hdfs://172.29.5.41:9000/ids/dataset/CIC_IDS2018_Clean_Balanced_25Percent.csv",
    header=True,
    inferSchema=False
)

feature_cols = [
    'Destination Port', 'Flow Duration', 'Total Fwd Packets',
    'Total Backward Packets', 'Fwd Packet Length Max', 'Fwd Packet Length Min',
    'Fwd Packet Length Mean', 'Bwd Packet Length Max', 'Bwd Packet Length Min',
    'Bwd Packet Length Mean', 'Flow Bytes/s', 'Flow Packets/s',
    'Flow IAT Mean', 'Flow IAT Max', 'Flow IAT Min',
    'Fwd IAT Mean', 'Fwd IAT Max', 'Fwd IAT Min',
    'Bwd IAT Mean', 'Bwd IAT Max', 'Bwd IAT Min',
    'Fwd Header Length', 'Bwd Header Length'
]

for c in feature_cols:
    df = df.withColumn(c, col(c).cast(DoubleType()))
df = df.withColumn("Label_Binary", col("Label_Binary").cast(IntegerType()))

df_clean = df.na.drop()
print("Rows after cleaning: " + str(df_clean.count()))

assembler = VectorAssembler(
    inputCols=feature_cols,
    outputCol="features",
    handleInvalid="skip"
)
data = assembler.transform(df_clean).select("features", "Label_Binary")

train_data, test_data = data.randomSplit([0.8, 0.2], seed=42)
print("Training rows: " + str(train_data.count()))
print("Test rows: " + str(test_data.count()))

rf = RandomForestClassifier(
    labelCol="Label_Binary",
    featuresCol="features",
    numTrees=100,
    maxDepth=15,
    seed=42
)

print("Starting distributed training...")
rf_model = rf.fit(train_data)
print("Training completed!")

predictions = rf_model.transform(test_data)

evaluator_acc = MulticlassClassificationEvaluator(
    labelCol="Label_Binary", predictionCol="prediction", metricName="accuracy")
evaluator_f1 = MulticlassClassificationEvaluator(
    labelCol="Label_Binary", predictionCol="prediction", metricName="f1")

accuracy = evaluator_acc.evaluate(predictions)
f1_score = evaluator_f1.evaluate(predictions)

print("=" * 50)
print("Test Accuracy: {:.2f}%".format(accuracy * 100))
print("Test F1-Score: {:.2f}%".format(f1_score * 100))
print("=" * 50)

# Save model to HDFS
rf_model.write().overwrite().save("hdfs://172.29.5.41:9000/ids/models/distributed_rf_model")
print("Model saved to HDFS")

spark.stop()
