import time
import json
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, trim
from pyspark.ml.feature import VectorAssembler, StringIndexer
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.evaluation import MulticlassClassificationEvaluator

spark = SparkSession.builder \
    .appName("Experiment1-RF-25pct") \
    .master("spark://10.0.0.15:7077") \
    .getOrCreate()

spark.sparkContext.setLogLevel("WARN")
print("Master: " + spark.sparkContext.master)
print("Default parallelism: " + str(spark.sparkContext.defaultParallelism))

# ---------- 1) Read clean data + sample 25% ----------
df_full = spark.read.parquet("hdfs://10.0.0.15:9000/ids/dataset_full_clean_v2.parquet")
print("Full clean rows: " + str(df_full.count()))

df = df_full.sample(withReplacement=False, fraction=0.25, seed=42)
df = df.repartition(6)
df.cache()
print("Sample (25%) rows: " + str(df.count()))
df.groupBy("Label").count().orderBy("count", ascending=False).show(20, truncate=False)

# ---------- 2) Feature columns (drop non-numeric / identifier columns) ----------
drop_cols = ["Timestamp", "Label", "Label_Binary"]
feature_cols = [c for c in df.columns if c not in drop_cols]
print("Number of features: " + str(len(feature_cols)))

assembler = VectorAssembler(inputCols=feature_cols, outputCol="features", handleInvalid="skip")

# ---------- 3) STAGE 1: Binary (Benign vs Attack) with class weights ----------
print("=" * 70)
print("STAGE 1 — Binary Classification (Random Forest)")
print("=" * 70)

data_bin = assembler.transform(df).select("features", "Label_Binary", "Label")

# Class weights to handle imbalance
counts = data_bin.groupBy("Label_Binary").count().collect()
total = sum([r["count"] for r in counts])
class_weights = {r["Label_Binary"]: total / (2.0 * r["count"]) for r in counts}
print("Class weights: " + str(class_weights))

from pyspark.sql.functions import udf
from pyspark.sql.types import DoubleType
weight_udf = udf(lambda l: float(class_weights[l]), DoubleType())
data_bin = data_bin.withColumn("weight", weight_udf(col("Label_Binary")))

train1, test1 = data_bin.randomSplit([0.8, 0.2], seed=42)
train1.cache(); test1.cache()
print("Stage1 Train: " + str(train1.count()) + " | Test: " + str(test1.count()))

rf1 = RandomForestClassifier(labelCol="Label_Binary", featuresCol="features",
                              weightCol="weight", numTrees=100, maxDepth=20, seed=42)

t0 = time.time()
model1 = rf1.fit(train1)
train_time_1 = time.time() - t0

t0 = time.time()
pred1 = model1.transform(test1)
pred1.cache()
n_test1 = pred1.count()
inference_time_1 = time.time() - t0
throughput_1 = n_test1 / inference_time_1

acc1 = MulticlassClassificationEvaluator(labelCol="Label_Binary", predictionCol="prediction", metricName="accuracy").evaluate(pred1)
f1_1 = MulticlassClassificationEvaluator(labelCol="Label_Binary", predictionCol="prediction", metricName="f1").evaluate(pred1)
prec1 = MulticlassClassificationEvaluator(labelCol="Label_Binary", predictionCol="prediction", metricName="weightedPrecision").evaluate(pred1)
rec1 = MulticlassClassificationEvaluator(labelCol="Label_Binary", predictionCol="prediction", metricName="weightedRecall").evaluate(pred1)

tp1 = pred1.filter((col("Label_Binary")==1) & (col("prediction")==1)).count()
tn1 = pred1.filter((col("Label_Binary")==0) & (col("prediction")==0)).count()
fp1 = pred1.filter((col("Label_Binary")==0) & (col("prediction")==1)).count()
fn1 = pred1.filter((col("Label_Binary")==1) & (col("prediction")==0)).count()

print("Stage1 -> Acc={:.4f} F1={:.4f} Prec={:.4f} Rec={:.4f}".format(acc1, f1_1, prec1, rec1))
print("Stage1 -> TP={} TN={} FP={} FN={}".format(tp1, tn1, fp1, fn1))
print("Stage1 -> Train time: {:.2f}s | Inference time: {:.2f}s | Throughput: {:.1f} rows/s".format(
    train_time_1, inference_time_1, throughput_1))

# ---------- 4) STAGE 2: Multiclass (attack type) on attack-only samples ----------
print("=" * 70)
print("STAGE 2 — Multiclass Classification (Random Forest) — Attack types only")
print("=" * 70)

df_attacks = df.filter(col("Label_Binary") == 1)
print("Attack-only rows: " + str(df_attacks.count()))
df_attacks.groupBy("Label").count().orderBy("count", ascending=False).show(20, truncate=False)

indexer = StringIndexer(inputCol="Label", outputCol="LabelIndex", handleInvalid="keep")
idx_model = indexer.fit(df_attacks)
df_attacks_idx = idx_model.transform(df_attacks)
labels_list = idx_model.labels
print("Attack classes: " + str(labels_list))

data_multi = assembler.transform(df_attacks_idx).select("features", "LabelIndex", "Label")

# Class weights for multiclass
counts_m = data_multi.groupBy("LabelIndex").count().collect()
total_m = sum([r["count"] for r in counts_m])
n_classes = len(counts_m)
class_weights_m = {r["LabelIndex"]: total_m / (n_classes * r["count"]) for r in counts_m}

weight_udf_m = udf(lambda l: float(class_weights_m[l]), DoubleType())
data_multi = data_multi.withColumn("weight", weight_udf_m(col("LabelIndex")))

train2, test2 = data_multi.randomSplit([0.8, 0.2], seed=42)
train2.cache(); test2.cache()
print("Stage2 Train: " + str(train2.count()) + " | Test: " + str(test2.count()))

rf2 = RandomForestClassifier(labelCol="LabelIndex", featuresCol="features",
                              weightCol="weight", numTrees=100, maxDepth=20, seed=42)

t0 = time.time()
model2 = rf2.fit(train2)
train_time_2 = time.time() - t0

t0 = time.time()
pred2 = model2.transform(test2)
pred2.cache()
n_test2 = pred2.count()
inference_time_2 = time.time() - t0
throughput_2 = n_test2 / inference_time_2

acc2 = MulticlassClassificationEvaluator(labelCol="LabelIndex", predictionCol="prediction", metricName="accuracy").evaluate(pred2)
f1_2 = MulticlassClassificationEvaluator(labelCol="LabelIndex", predictionCol="prediction", metricName="f1").evaluate(pred2)
prec2 = MulticlassClassificationEvaluator(labelCol="LabelIndex", predictionCol="prediction", metricName="weightedPrecision").evaluate(pred2)
rec2 = MulticlassClassificationEvaluator(labelCol="LabelIndex", predictionCol="prediction", metricName="weightedRecall").evaluate(pred2)

print("Stage2 -> Acc={:.4f} F1={:.4f} Prec={:.4f} Rec={:.4f}".format(acc2, f1_2, prec2, rec2))
print("Stage2 -> Train time: {:.2f}s | Inference time: {:.2f}s | Throughput: {:.1f} rows/s".format(
    train_time_2, inference_time_2, throughput_2))

# Per-class confusion for stage 2 (One-vs-Rest)
print("Per-class breakdown (Stage 2):")
for idx, label_name in enumerate(labels_list):
    tp = pred2.filter((col("LabelIndex")==idx) & (col("prediction")==idx)).count()
    fp = pred2.filter((col("LabelIndex")!=idx) & (col("prediction")==idx)).count()
    fn = pred2.filter((col("LabelIndex")==idx) & (col("prediction")!=idx)).count()
    print("  {} -> TP={} FP={} FN={}".format(label_name, tp, fp, fn))

# ---------- 5) Save results ----------
results = {
    "experiment": "Experiment1_RF_25pct_3nodes",
    "sample_fraction": 0.25,
    "num_features": len(feature_cols),
    "stage1_binary": {
        "accuracy": acc1, "f1": f1_1, "precision": prec1, "recall": rec1,
        "TP": tp1, "TN": tn1, "FP": fp1, "FN": fn1,
        "train_time_sec": train_time_1, "inference_time_sec": inference_time_1,
        "throughput_rows_per_sec": throughput_1
    },
    "stage2_multiclass": {
        "accuracy": acc2, "f1": f1_2, "precision": prec2, "recall": rec2,
        "attack_classes": labels_list,
        "train_time_sec": train_time_2, "inference_time_sec": inference_time_2,
        "throughput_rows_per_sec": throughput_2
    }
}

with open("/opt/spark-apps/experiment1_results.json", "w") as f:
    json.dump(results, f, indent=2)

print("=" * 70)
print("RESULTS SAVED to /opt/spark-apps/experiment1_results.json")
print("=" * 70)

# Save models to HDFS
model1.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp1_rf_stage1")
model2.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp1_rf_stage2")
idx_model.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp1_indexer")

spark.stop()
