from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, trim, isnan
from pyspark.sql.types import DoubleType

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

spark.sparkContext.setLogLevel("WARN")

files = ["02-14-2018.csv", "02-15-2018.csv", "02-16-2018.csv",
         "02-20-2018.csv", "02-21-2018.csv", "02-22-2018.csv",
         "02-23-2018.csv", "02-28-2018.csv", "03-01-2018.csv",
         "03-02-2018.csv"]

extra_cols_in_0220 = ["Flow ID", "Src IP", "Src Port", "Dst IP"]

dfs = []
for f in files:
    path = "hdfs://10.0.0.15:9000/ids/raw_dataset/" + f
    d = spark.read.csv(path, header=True, inferSchema=True)
    cols_to_drop = [c for c in extra_cols_in_0220 if c in d.columns]
    if cols_to_drop:
        print(f + " -> dropping extra columns: " + str(cols_to_drop))
        d = d.drop(*cols_to_drop)
    dfs.append(d)

df = dfs[0]
for d in dfs[1:]:
    df = df.unionByName(d, allowMissingColumns=False)

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

# Remove duplicated header rows leaked inside the body
df = df.filter(col("Label") != "Label")

# ---- THE FIX: force-cast every feature column to Double, regardless of inferred type ----
exclude_cols = ["Label", "Timestamp"]
feature_cols = [c for c in df.columns if c not in exclude_cols]

print("Force-casting " + str(len(feature_cols)) + " feature columns to Double...")
for c in feature_cols:
    df = df.withColumn(c, col(c).cast(DoubleType()))

rows_before = df.count()

# Now remove NULL, NaN, and Infinity robustly using explicit equality (not isin)
for c in feature_cols:
    df = df.filter(
        col(c).isNotNull() &
        (~isnan(col(c))) &
        (col(c) != float("inf")) &
        (col(c) != float("-inf"))
    )

rows_after = df.count()
print("Rows before cleaning: " + str(rows_before))
print("Rows after cleaning : " + str(rows_after))
print("Rows removed        : " + str(rows_before - rows_after))

# Exclude any leaked numeric-looking labels
df = df.filter(~col("Label").rlike("^[0-9]"))

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

# ---- Verification step: confirm ZERO infinities remain in the two problem columns ----
check_cols = ["Flow Byts/s", "Flow Pkts/s"]
for c in check_cols:
    if c in df.columns:
        inf_count = df.filter((col(c) == float("inf")) | (col(c) == float("-inf"))).count()
        nan_count = df.filter(isnan(col(c))).count()
        print("VERIFY " + c + " -> remaining Infinity: " + str(inf_count) + " | remaining NaN: " + str(nan_count))

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

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

spark.stop()
