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

spark = SparkSession.builder \
    .appName("Experiment3-RF-100pct-Cluster") \
    .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 FULL clean data (100%, no sampling) ----------
df = spark.read.parquet("hdfs://10.0.0.15:9000/ids/dataset_full_clean_v2.parquet")
df = df.repartition(12)  # more partitions for larger data volume
df.cache()
print("Full dataset (100%) rows: " + str(df.count()))
df.groupBy("Label").count().orderBy("count", ascending=False).show(20, truncate=False)

# ---------- 2) Feature columns (identical to all previous experiments) ----------
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) — FULL DATASET")
print("=" * 70)

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

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))

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))

# Checkpoint immediately
with open("exp3_stage1_checkpoint.json", "w") as f:
    json.dump({
        "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
        }
    }, f, indent=2)
print("Stage 1 checkpoint saved.")

model1.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp3_rf_stage1")

# ---------- 4) STAGE 2: Multiclass (attack type) on attack-only samples ----------
print("=" * 70)
print("STAGE 2 — Multiclass Classification (Random Forest) — FULL DATASET")
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")

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))

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))

results = {
    "experiment": "Experiment3_RF_100pct_3nodes_cluster",
    "sample_fraction": 1.0,
    "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/experiment3_rf_100pct_results.json", "w") as f:
    json.dump(results, f, indent=2)

print("RESULTS SAVED to /opt/spark-apps/experiment3_rf_100pct_results.json")

model2.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp3_rf_stage2")
idx_model.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp3_rf_indexer")

spark.stop()
