import time
import json
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, trim, lit, udf
from pyspark.sql.types import StringType, DoubleType
from pyspark.ml.feature import VectorAssembler, StringIndexerModel
from pyspark.ml.classification import GBTClassificationModel, OneVsRestModel
from pyspark.mllib.evaluation import MulticlassMetrics

spark = SparkSession.builder \
    .appName("Experiment1-GBT-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 (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) SAME test split as training ----------
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: " + str(n_test_total))

# ---------- 3) Load trained GBT models from HDFS ----------
print("Loading Stage 1 (GBT Binary) model from HDFS...")
model_stage1 = GBTClassificationModel.load("hdfs://10.0.0.15:9000/ids/models/exp1_gbt_stage1")

print("Loading Stage 2 (OneVsRest+GBT Multiclass) model from HDFS...")
model_stage2 = OneVsRestModel.load("hdfs://10.0.0.15:9000/ids/models/exp1_gbt_stage2")

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

# ---------- 4) STAGE 1 inference ----------
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: " + str(stage1_inference_time) + "s for " + str(n1) + " rows")

# Rename to avoid column collision with Stage 2
pred_stage1 = pred_stage1.withColumnRenamed("prediction", "stage1_prediction") \
                          .withColumnRenamed("rawPrediction", "stage1_rawPrediction") \
                          .withColumnRenamed("probability", "stage1_probability")

predicted_attacks = pred_stage1.filter(col("stage1_prediction") == 1.0)
predicted_benign  = pred_stage1.filter(col("stage1_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 — GBT")
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))

# ---------- 5) Build final predicted label ----------
predicted_benign_final = predicted_benign.withColumn("final_predicted_label", lit("Benign"))

if pred_stage2 is not None:
    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

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

# ---------- 6) Unified 15-class evaluation ----------
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))

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

label_to_idx_udf = udf(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"))
)

unified_results = unified_results.filter((col("true_idx") >= 0) & (col("pred_idx") >= 0))
unified_results.cache()

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 — GBT FULL HIERARCHICAL PIPELINE")
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))

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)

final_results = {
    "experiment": "Experiment1_GBT_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_gbt_unified_results.json", "w") as f:
    json.dump(final_results, f, indent=2)

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

spark.stop()
