import time
import json
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, trim, lit
from pyspark.ml.feature import VectorAssembler, StringIndexerModel
from pyspark.ml.classification import RandomForestClassificationModel
from pyspark.mllib.evaluation import MulticlassMetrics

spark = SparkSession.builder \
    .appName("Experiment1-Unified-Pipeline-Eval") \
    .master("spark://10.0.0.15:7077") \
    .getOrCreate()

spark.sparkContext.setLogLevel("WARN")
print("Master: " + spark.sparkContext.master)

# ---------- 1) Load the SAME 25% sample used in training (same seed = same split) ----------
df_full = spark.read.parquet("hdfs://10.0.0.15:9000/ids/dataset_full_clean_v2.parquet")
df = df_full.sample(withReplacement=False, fraction=0.25, seed=42)
df = df.repartition(6)
df.cache()
print("Sample (25%) rows: " + str(df.count()))

drop_cols = ["Timestamp", "Label", "Label_Binary"]
feature_cols = [c for c in df.columns if c not in drop_cols]
assembler = VectorAssembler(inputCols=feature_cols, outputCol="features", handleInvalid="skip")

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

# ---------- 2) Recreate the EXACT SAME test split used during training ----------
# We split the FULL assembled data the same way Stage 1 was split originally
train_full, test_full = df_assembled.randomSplit([0.8, 0.2], seed=42)
test_full.cache()

n_test_total = test_full.count()
print("Unified test set (original 20% holdout, includes Benign + Attacks): " + str(n_test_total))

# ---------- 3) Load trained models from HDFS (no retraining) ----------
print("Loading Stage 1 (Binary) model from HDFS...")
model_stage1 = RandomForestClassificationModel.load("hdfs://10.0.0.15:9000/ids/models/exp1_rf_stage1")

print("Loading Stage 2 (Multiclass) model from HDFS...")
model_stage2 = RandomForestClassificationModel.load("hdfs://10.0.0.15:9000/ids/models/exp1_rf_stage2")

print("Loading StringIndexer model from HDFS...")
indexer_model = StringIndexerModel.load("hdfs://10.0.0.15:9000/ids/models/exp1_indexer")
attack_labels = indexer_model.labels
print("Attack classes (Stage 2 index mapping): " + str(attack_labels))

# ---------- 4) STAGE 1 inference on unified test set ----------
t0 = time.time()
pred_stage1 = model_stage1.transform(test_full)
pred_stage1.cache()
n1 = pred_stage1.count()
stage1_inference_time = time.time() - t0
print("Stage 1 inference on unified test set: " + str(stage1_inference_time) + "s for " + str(n1) + " rows")

# Rename Stage 1's "prediction" and other ML columns BEFORE passing to Stage 2
# to avoid column name collision (Stage 2 also creates "prediction", "rawPrediction", "probability")
pred_stage1 = pred_stage1.withColumnRenamed("prediction", "stage1_prediction") \
                          .withColumnRenamed("rawPrediction", "stage1_rawPrediction") \
                          .withColumnRenamed("probability", "stage1_probability")

# ---------- 5) Route: rows predicted as Attack (stage1_prediction == 1) go to Stage 2 ----------
predicted_attacks = pred_stage1.filter(col("stage1_prediction") == 1.0)
predicted_benign  = pred_stage1.filter(col("stage1_prediction") == 0.0)

# ---------- 5) Route: rows predicted as Attack (prediction == 1) go to Stage 2 ----------
predicted_attacks = pred_stage1.filter(col("prediction") == 1.0)
predicted_benign  = pred_stage1.filter(col("prediction") == 0.0)

n_predicted_attacks = predicted_attacks.count()
n_predicted_benign  = predicted_benign.count()
print("Routed to Stage 2 (predicted Attack): " + str(n_predicted_attacks))
print("Stopped at Stage 1 (predicted Benign): " + str(n_predicted_benign))

t0 = time.time()
if n_predicted_attacks > 0:
    pred_stage2 = model_stage2.transform(predicted_attacks)
    pred_stage2.cache()
    n2 = pred_stage2.count()
else:
    pred_stage2 = None
    n2 = 0
stage2_inference_time = time.time() - t0
print("Stage 2 inference: " + str(stage2_inference_time) + "s for " + str(n2) + " rows")

total_inference_time = stage1_inference_time + stage2_inference_time
unified_throughput = n_test_total / total_inference_time
print("=" * 70)
print("UNIFIED (END-TO-END) INFERENCE PERFORMANCE")
print("=" * 70)
print("Total rows through pipeline : " + str(n_test_total))
print("Total inference time        : {:.2f}s".format(total_inference_time))
print("Unified throughput          : {:.1f} rows/sec".format(unified_throughput))
print("Latency per row              : {:.4f} ms/row".format((total_inference_time / n_test_total) * 1000))

# ---------- 6) Build FINAL predicted label for EVERY row in the unified test set ----------
# Rows predicted Benign at Stage 1 -> final label = "Benign"
predicted_benign_final = predicted_benign.withColumn("final_predicted_label", lit("Benign"))

# Rows predicted Attack at Stage 1 -> final label = Stage 2 attack type prediction
if pred_stage2 is not None:
    # Map numeric Stage2 prediction back to string label
    from pyspark.sql.types import StringType
    from pyspark.sql.functions import udf

    labels_broadcast = attack_labels
    def idx_to_label(idx):
        idx = int(idx)
        if idx < len(labels_broadcast):
            return labels_broadcast[idx]
        return "Unknown"

    idx_to_label_udf = udf(idx_to_label, StringType())
    predicted_attack_final = pred_stage2.withColumn(
        "final_predicted_label", idx_to_label_udf(col("prediction"))
    )
else:
    predicted_attack_final = None

# ---------- 7) Combine both branches into one unified result set ----------
cols_needed = ["Label", "final_predicted_label"]

unified_benign = predicted_benign_final.select(*cols_needed)

if predicted_attack_final is not None:
    unified_attacks = predicted_attack_final.select(*cols_needed)
    unified_results = unified_benign.unionByName(unified_attacks)
else:
    unified_results = unified_benign

unified_results.cache()
n_unified = unified_results.count()
print("Unified results rows (should equal " + str(n_test_total) + "): " + str(n_unified))

# ---------- 8) Build a single label index space: all 15 classes (Benign + 14 attack types) ----------
all_classes = ["Benign"] + attack_labels
class_to_idx = {name: i for i, name in enumerate(all_classes)}
print("Unified class space (15 classes): " + str(all_classes))

from pyspark.sql.types import DoubleType
from pyspark.sql.functions import udf as udf2

def label_to_idx(name):
    return float(class_to_idx.get(name, -1))

label_to_idx_udf = udf2(label_to_idx, DoubleType())

unified_results = unified_results.withColumn(
    "true_idx", label_to_idx_udf(col("Label"))
).withColumn(
    "pred_idx", label_to_idx_udf(col("final_predicted_label"))
)

# Drop any row where mapping failed (label not recognized) — should not normally happen
unified_results = unified_results.filter((col("true_idx") >= 0) & (col("pred_idx") >= 0))
unified_results.cache()

# ---------- 9) Compute UNIFIED metrics for the full hierarchical pipeline ----------
pred_and_labels = unified_results.select("pred_idx", "true_idx").rdd.map(
    lambda row: (row[0], row[1])
)

metrics = MulticlassMetrics(pred_and_labels)

overall_accuracy = metrics.accuracy
overall_weighted_precision = metrics.weightedPrecision
overall_weighted_recall = metrics.weightedRecall
overall_weighted_f1 = metrics.weightedFMeasure()

print("=" * 70)
print("UNIFIED SYSTEM METRICS — FULL HIERARCHICAL PIPELINE (Stage1 + Stage2)")
print("=" * 70)
print("Overall Accuracy          : {:.4f}".format(overall_accuracy))
print("Overall Weighted F1       : {:.4f}".format(overall_weighted_f1))
print("Overall Weighted Precision: {:.4f}".format(overall_weighted_precision))
print("Overall Weighted Recall   : {:.4f}".format(overall_weighted_recall))

# ---------- 10) Per-class TP / FP / FN across the UNIFIED 15-class space ----------
print("-" * 70)
print("{:<28}{:<10}{:<10}{:<10}{:<10}{:<10}".format(
    "Class", "Precision", "Recall", "F1", "TP", "FP") + "FN")
print("-" * 70)

per_class_results = []
for cname, cidx in class_to_idx.items():
    precision = metrics.precision(float(cidx))
    recall    = metrics.recall(float(cidx))
    f1        = metrics.fMeasure(float(cidx))

    tp = unified_results.filter((col("true_idx") == cidx) & (col("pred_idx") == cidx)).count()
    fp = unified_results.filter((col("true_idx") != cidx) & (col("pred_idx") == cidx)).count()
    fn = unified_results.filter((col("true_idx") == cidx) & (col("pred_idx") != cidx)).count()

    print("{:<28}{:<10.4f}{:<10.4f}{:<10.4f}{:<10}{:<10}{}".format(
        cname, precision, recall, f1, tp, fp, fn))

    per_class_results.append({
        "class": cname, "precision": precision, "recall": recall,
        "f1": f1, "tp": tp, "fp": fp, "fn": fn
    })

print("=" * 70)

# ---------- 11) Save everything to JSON ----------
final_results = {
    "experiment": "Experiment1_RF_25pct_UNIFIED_pipeline_eval",
    "test_set_size": n_test_total,
    "routed_to_stage2": n_predicted_attacks,
    "stopped_at_stage1": n_predicted_benign,
    "timing": {
        "stage1_inference_time_sec": stage1_inference_time,
        "stage2_inference_time_sec": stage2_inference_time,
        "total_inference_time_sec": total_inference_time,
        "unified_throughput_rows_per_sec": unified_throughput,
        "latency_ms_per_row": (total_inference_time / n_test_total) * 1000
    },
    "unified_metrics": {
        "accuracy": overall_accuracy,
        "weighted_f1": overall_weighted_f1,
        "weighted_precision": overall_weighted_precision,
        "weighted_recall": overall_weighted_recall
    },
    "per_class": per_class_results
}

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

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

spark.stop()
