Kafka (real-time clicks) 
        ↓
Spark Structured Streaming 
        ↓
Real-time Feature Computation (Feature Pipelines) 
        ↓
Online Model Inference (Batch or Streaming models) 
        ↓
→ Low-latency prediction (<100ms) 
→ Write prediction back to Kafka + Redis + Dashboard

+----------------+      +--------------------- +      +-----------------+
|   Kafka        | ---> | Spark Structured     | ---> | Redis (serving) |
| clicks/events  |      | Streaming + MLlib    |      | Kafka (results) |
+----------------+      | or Spark + PMML/MLeap|      +-----------------+
                               +---------------------+
                                             ↓
                                     Live Dashboard (Streamlit)

# Run in Databricks or Colab
from pyspark.sql import SparkSession
from pyspark.ml.feature import VectorAssembler, StringIndexer
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml import Pipeline

spark = SparkSession.builder.getOrCreate()

# Public credit card fraud dataset (1.8M rows)
df = spark.read.csv("s3://spark-ml-data/creditcard.csv", header=True, inferSchema=True)

# Features V1..V28 + Amount are already PCA-transformed
feature_cols = [f"V{i}" for i in range(1,29)] + ["Amount"]

assembler = VectorAssembler(inputCols=feature_cols, outputCol="features")
label_indexer = StringIndexer(inputCol="Class", outputCol="label")

rf = RandomForestClassifier(featuresCol="features", labelCol="label", numTrees=100, maxDepth=10)

pipeline = Pipeline(stages=[assembler, label_indexer, rf])

# Train on full data (in production you would use train/test split)
model = pipeline.fit(df)

# Save model – works with both batch and streaming
model.write().overwrite().save("/tmp/models/fraud_rf_model")
print("Model trained and saved!")

from pyspark.ml import PipelineModel
from pyspark.sql.functions import struct, col

# Load the trained model (once at driver startup)
model = PipelineModel.load("/tmp/models/fraud_rf_model")

# Simulate real-time transaction stream
schema = "Time FLOAT, V1 FLOAT, V2 FLOAT, ..., V28 FLOAT, Amount FLOAT"
# In real life: read from Kafka with proper schema

def realtime_inference():
    streaming_df = spark \
        .readStream \
        .format("socket") \
        .option("host", "localhost") \
        .option("port", 9999) \
        .load()

    # Parse CSV → struct (in production use Kafka + schema registry)
    from pyspark.sql.functions import from_csv
    options = {"header": "false", "inferSchema": "true"}
    parsed = streaming_df.select(from_csv(col("value"), schema, options).alias("data")).select("data.*")

    # Real-time prediction
    predictions = model.transform(parsed)

    # Select only needed fields for low latency
    result = predictions.select(
        "Amount",
        "prediction",
        col("probability")[1].alias("fraud_probability"),
        current_timestamp().alias("prediction_time")
    )

    # Sink 1: Console (debug)
    # Sink 2: Redis (ultra-low latency serving)
    # Sink 3: Kafka (downstream systems)

    query = result \
        .writeStream \
        .outputMode("append") \
        .format("console") \
        .start()

    # Sink to Redis (using foreachBatch + Redis-py)
    def write_to_redis(batch_df, batch_id):
        pdf = batch_df.select("prediction", "fraud_probability", "Amount").toPandas()
        import redis
        r = redis.Redis(host='localhost', port=6379, db=0)
        for _, row in pdf.iterrows():
            key = f"tx:{batch_id}_{_}"
            r.hset(key, mapping=row.to_dict())
            r.expire(key, 3600)  # keep 1 hour

    redis_query = result \
        .writeStream \
        .foreachBatch(write_to_redis) \
        .start()

    return query, redis_query

# q1, q2 = realtime_inference()
# q1.awaitTermination()

# In your local environment or Databricks
pip install mleap onnxruntime

# Convert Spark model to MLeap bundle
from mleap.pyspark.spark_support import SimpleSparkSerializer

model.stages[-1].serializeToBundle("jar:file:/tmp/fraud_model.zip", model.transform(df.limit(1)))

# Now you can load this model in a Flask/FastAPI service with <10ms latency!

# fast_inference_service.py
from mleap.pyspark.spark_support import SimpleSparkSerializer
import pandas as pd

# Load once at startup
model = SimpleSparkSerializer().deserializeFromBundle("jar:file:/tmp/fraud_model.zip")

def predict(features: list) -> dict:
    df = pd.DataFrame([features], columns=[f"V{i}" for i in range(1,29)] + ["Amount"])
    prediction = model.transform(df)
    return {
        "fraud_score": prediction["probability"].iloc[0][1],
        "is_fraud": int(prediction["prediction"].iloc[0])
    }

# Example with River (Python online ML library) inside foreachBatch
def online_training_and_inference(batch_df, batch_id):
    import river
    from river import linear_model, preprocessing, metrics

    model = (preprocessing.StandardScaler() |
             linear_model.LogisticRegression())

    metric = metrics.Accuracy()

    for _, row in batch_df.toPandas().iterrows():
        features = row[[f"V{i}" for i in range(1,29)] + ["Amount"]].values
        y_true = row["Class"]

        y_pred = model.predict_one(dict(enumerate(features)))
        model.learn_one(dict(enumerate(features)), y_true)
        metric.update(y_true, y_pred)

    print(f"Batch {batch_id} accuracy: {metric.get()}")