from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, trim

spark = SparkSession.builder \
    .appName("Clean-Full-Dataset") \
    .master("spark://10.0.0.15:7077") \
    .getOrCreate()

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

df = spark.read.csv("hdfs://10.0.0.15:9000/ids/raw_dataset/*.csv",
                     header=True, inferSchema=True)

print("Raw rows: " + str(df.count()))
print("Columns: " + str(df.columns))

numeric_cols = [f.name for f in df.schema.fields
                if str(f.dataType) in ("DoubleType", "FloatType")]

for c in numeric_cols:
    df = df.filter(~col(c).isin(float("inf"), float("-inf")))

df = df.na.drop()

df = df.withColumn("Label_Binary",
        when(trim(col("Label")) == "Benign", 0).otherwise(1))

df.write.mode("overwrite").parquet("hdfs://10.0.0.15:9000/ids/dataset_full_clean.parquet")

print("Total clean rows: " + str(df.count()))
df.groupBy("Label").count().orderBy("count", ascending=False).show(50, truncate=False)

spark.stop()
