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 GBTClassifier, OneVsRest
from pyspark.ml.evaluation import MulticlassClassificationEvaluator

spark = SparkSession.builder \
    .appName("Experiment1-GBT-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 + SAME 25% sample (same seed as RF experiment) ----------
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 (identical to RF experiment) ----------
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 — GBT direct ----------
print("=" * 70)
print("STAGE 1 — Binary Classification (GBT)")
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")))

# SAME split (same seed) as RF experiment — critical for fair comparison
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()))

# GBTClassifier supports weightCol natively — same maxDepth as RF (20 -> GBT max is 30, but
# we keep depth comparable; maxIter roughly analogous to numTrees but GBT is sequential so
# we use a smaller iteration count typical for GBT to keep training time reasonable)
gbt1 = GBTClassifier(labelCol="Label_Binary", featuresCol="features",
                      weightCol="weight", maxIter=50, maxDepth=10, seed=42)

t0 = time.time()
model1 = gbt1.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 via OneVsRest(GBT) — attack types only ----------
print("=" * 70)
print("STAGE 2 — Multiclass Classification (OneVsRest + GBT) — 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")

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

# NOTE: OneVsRest does not support weightCol directly pre-3.x uniformly across all versions,
# so class imbalance here is handled implicitly by GBT's boosting mechanism itself,
# which naturally focuses more on hard/minority examples each iteration.
base_gbt = GBTClassifier(labelCol="LabelIndex", featuresCol="features",
                          maxIter=20, maxDepth=10, seed=42)

ovr = OneVsRest(labelCol="LabelIndex", featuresCol="features", classifier=base_gbt)

t0 = time.time()
model2 = ovr.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))

# ---------- 5) Save results ----------
results = {
    "experiment": "Experiment1_GBT_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_gbt_results.json", "w") as f:
    json.dump(results, f, indent=2)

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

# Save models to HDFS
model1.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp1_gbt_stage1")
model2.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp1_gbt_stage2")
idx_model.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp1_gbt_indexer")

spark.stop()
