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

spark = SparkSession.builder \
    .appName("Experiment1-MLP-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/GBT experiments) ----------
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/GBT 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_raw", handleInvalid="skip")

# StandardScaler is applied ONLY for MLP (not for RF/GBT), because gradient-based
# optimization (L-BFGS) is highly sensitive to feature scale differences,
# unlike tree-based splits used in RF/GBT which are scale-invariant.
scaler = StandardScaler(inputCol="features_raw", outputCol="features",
                         withMean=True, withStd=True)

# ---------- 3) STAGE 1: Binary (Benign vs Attack) — MLP with manual oversampling ----------
print("=" * 70)
print("STAGE 1 — Binary Classification (MLP)")
print("=" * 70)

assembled_bin = assembler.transform(df).select("features_raw", "Label_Binary", "Label")

# Fit scaler on the FULL assembled data before splitting (standard practice: scaler
# statistics should reflect the overall feature distribution)
scaler_model_stage1 = scaler.fit(assembled_bin)
data_bin = scaler_model_stage1.transform(assembled_bin).select("features", "Label_Binary", "Label")

# SAME split (same seed) as RF/GBT experiments — critical for fair comparison
train1_raw, test1 = data_bin.randomSplit([0.8, 0.2], seed=42)
train1_raw.cache(); test1.cache()

# Manual oversampling of minority class (Attack=1) ONLY on training set
counts = train1_raw.groupBy("Label_Binary").count().collect()
count_map = {r["Label_Binary"]: r["count"] for r in counts}
print("Stage1 raw train distribution: " + str(count_map))

majority_label = max(count_map, key=count_map.get)
minority_label = min(count_map, key=count_map.get)
ratio = count_map[majority_label] / count_map[minority_label]
print("Oversampling ratio for minority class (" + str(minority_label) + "): " + str(ratio))

majority_df = train1_raw.filter(col("Label_Binary") == majority_label)
minority_df = train1_raw.filter(col("Label_Binary") == minority_label)

minority_oversampled = minority_df.sample(withReplacement=True, fraction=ratio, seed=42)

train1 = majority_df.unionByName(minority_oversampled)
train1.cache()
print("Stage1 Train (after oversampling): " + str(train1.count()) + " | Test: " + str(test1.count()))
train1.groupBy("Label_Binary").count().show()

# MLP layers: [input=78, hidden=40, hidden=20, output=2]
layers1 = [len(feature_cols), 40, 20, 2]
print("Stage1 MLP layers: " + str(layers1))

mlp1 = MultilayerPerceptronClassifier(labelCol="Label_Binary", featuresCol="features",
                                       layers=layers1, maxIter=60, blockSize=128, seed=42)

t0 = time.time()
model1 = mlp1.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 — MLP native multiclass support ----------
print("=" * 70)
print("STAGE 2 — Multiclass Classification (MLP) — 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
num_classes = len(labels_list)
print("Attack classes: " + str(labels_list))
print("Number of attack classes: " + str(num_classes))

assembled_multi = assembler.transform(df_attacks_idx).select("features_raw", "LabelIndex", "Label")

# Fit a SEPARATE scaler for Stage 2 (different data distribution: attacks only)
scaler_model_stage2 = scaler.fit(assembled_multi)
data_multi = scaler_model_stage2.transform(assembled_multi).select("features", "LabelIndex", "Label")

train2_raw, test2 = data_multi.randomSplit([0.8, 0.2], seed=42)
train2_raw.cache(); test2.cache()

# Manual oversampling for Stage 2: bring every minority class up toward the majority count
counts_m = train2_raw.groupBy("LabelIndex").count().collect()
count_map_m = {r["LabelIndex"]: r["count"] for r in counts_m}
print("Stage2 raw train distribution: " + str(count_map_m))

max_count = max(count_map_m.values())
balanced_parts = []
for lbl, cnt in count_map_m.items():
    subset = train2_raw.filter(col("LabelIndex") == lbl)
    if cnt < max_count:
        factor = max_count / cnt
        # Cap oversampling factor to avoid extreme duplication/training time for rare classes
        factor = min(factor, 20.0)
        subset = subset.sample(withReplacement=True, fraction=factor, seed=42)
    balanced_parts.append(subset)

train2 = balanced_parts[0]
for part in balanced_parts[1:]:
    train2 = train2.unionByName(part)

train2.cache()
print("Stage2 Train (after oversampling): " + str(train2.count()) + " | Test: " + str(test2.count()))
train2.groupBy("LabelIndex").count().orderBy("LabelIndex").show(20)

# MLP layers: [input=78, hidden=50, hidden=30, output=num_classes]
layers2 = [len(feature_cols), 50, 30, num_classes]
print("Stage2 MLP layers: " + str(layers2))

mlp2 = MultilayerPerceptronClassifier(labelCol="LabelIndex", featuresCol="features",
                                       layers=layers2, maxIter=60, blockSize=128, seed=42)

t0 = time.time()
model2 = mlp2.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_MLP_25pct_3nodes",
    "sample_fraction": 0.25,
    "num_features": len(feature_cols),
    "layers_stage1": layers1,
    "layers_stage2": layers2,
    "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_mlp_results.json", "w") as f:
    json.dump(results, f, indent=2)

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

# Save models to HDFS
model1.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp1_mlp_stage1")
model2.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp1_mlp_stage2")
idx_model.write().overwrite().save("hdfs://10.0.0.15:9000/ids/models/exp1_mlp_indexer")

spark.stop()
