# Databricks notebook example - Big Data Analytics
from pyspark.sql import SparkSession
from pyspark.sql.functions import *
from pyspark.sql.types import *
import mlflow
import mlflow.spark
from databricks import feature_store

# Initialize Spark session (automatic in Databricks)
spark = SparkSession.getActiveSession()

# Data ingestion from multiple sources
def load_and_process_data():
    # Load from Delta Lake
    sales_df = spark.read.table("sales.raw_orders")
    
    # Load from S3/ADLS
    customer_df = spark.read.option("header", "true")\
                      .csv("s3://data-lake/customers/")
    
    # Load streaming data
    streaming_df = spark.readStream\
                       .format("kafka")\
                       .option("kafka.bootstrap.servers", "broker:9092")\
                       .option("subscribe", "user_events")\
                       .load()
    
    return sales_df, customer_df, streaming_df

# Data transformation with Spark
def transform_sales_data(sales_df, customer_df):
    # Complex transformations
    enriched_df = sales_df.join(customer_df, "customer_id")\
                         .withColumn("order_month", date_format("order_date", "yyyy-MM"))\
                         .withColumn("revenue", col("quantity") * col("unit_price"))\
                         .withColumn("customer_lifetime_value", 
                                   sum("revenue").over(Window.partitionBy("customer_id")))
    
    # Feature engineering for ML
    features_df = enriched_df.groupBy("customer_id", "order_month")\
                            .agg(
                                count("*").alias("order_count"),
                                sum("revenue").alias("monthly_revenue"),
                                avg("revenue").alias("avg_order_value"),
                                max("order_date").alias("last_order_date")
                            )
    
    return features_df

# Machine Learning Pipeline
def ml_pipeline(features_df):
    from pyspark.ml import Pipeline
    from pyspark.ml.feature import VectorAssembler, StandardScaler
    from pyspark.ml.regression import RandomForestRegressor
    from pyspark.ml.evaluation import RegressionEvaluator
    
    # Feature preparation
    assembler = VectorAssembler(
        inputCols=["order_count", "monthly_revenue", "avg_order_value"],
        outputCol="features"
    )
    
    scaler = StandardScaler(
        inputCol="features",
        outputCol="scaled_features",
        withStd=True,
        withMean=True
    )
    
    # ML model
    rf = RandomForestRegressor(
        featuresCol="scaled_features",
        labelCol="monthly_revenue",
        numTrees=100
    )
    
    # Pipeline
    pipeline = Pipeline(stages=[assembler, scaler, rf])
    
    # Split data
    train_df, test_df = features_df.randomSplit([0.8, 0.2], seed=42)
    
    # Train model with MLflow tracking
    with mlflow.start_run(run_name="customer_revenue_prediction"):
        mlflow.log_param("num_trees", 100)
        mlflow.log_param("train_size", train_df.count())
        
        model = pipeline.fit(train_df)
        
        # Evaluate model
        predictions = model.transform(test_df)
        evaluator = RegressionEvaluator(
            labelCol="monthly_revenue",
            predictionCol="prediction",
            metricName="rmse"
        )
        
        rmse = evaluator.evaluate(predictions)
        mlflow.log_metric("rmse", rmse)
        
        # Log model
        mlflow.spark.log_model(model, "revenue_prediction_model")
        
    return model

# Execute pipeline
sales_df, customer_df, streaming_df = load_and_process_data()
features_df = transform_sales_data(sales_df, customer_df)
model = ml_pipeline(features_df)

# Save results to Delta Lake
features_df.write.format("delta").mode("overwrite").saveAsTable("features.customer_monthly")

print("Pipeline completed successfully!")

# Delta Lake advanced features
from delta.tables import DeltaTable
from pyspark.sql.functions import when

class DeltaLakeManager:
    def __init__(self, spark_session):
        self.spark = spark_session
    
    def create_delta_table(self, df, table_path, partition_cols=None):
        """Create Delta Lake table with optimizations"""
        writer = df.write.format("delta")
        
        if partition_cols:
            writer = writer.partitionBy(*partition_cols)
        
        writer.option("delta.autoOptimize.optimizeWrite", "true")\
              .option("delta.autoOptimize.autoCompact", "true")\
              .mode("overwrite")\
              .save(table_path)
        
        return DeltaTable.forPath(self.spark, table_path)
    
    def upsert_data(self, target_table_path, updates_df, merge_condition):
        """Perform UPSERT operation"""
        target_table = DeltaTable.forPath(self.spark, target_table_path)
        
        (target_table.alias("target")
         .merge(updates_df.alias("updates"), merge_condition)
         .whenMatchedUpdateAll()
         .whenNotMatchedInsertAll()
         .execute())
    
    def time_travel_query(self, table_path, version=None, timestamp=None):
        """Query historical data"""
        if version:
            return self.spark.read.format("delta").option("versionAsOf", version).load(table_path)
        elif timestamp:
            return self.spark.read.format("delta").option("timestampAsOf", timestamp).load(table_path)
    
    def optimize_table(self, table_path, z_order_cols=None):
        """Optimize Delta table performance"""
        if z_order_cols:
            self.spark.sql(f"OPTIMIZE delta.`{table_path}` ZORDER BY ({','.join(z_order_cols)})")
        else:
            self.spark.sql(f"OPTIMIZE delta.`{table_path}`")
    
    def vacuum_table(self, table_path, retention_hours=168):
        """Clean up old files"""
        self.spark.sql(f"VACUUM delta.`{table_path}` RETAIN {retention_hours} HOURS")

# Usage
delta_manager = DeltaLakeManager(spark)

# Create optimized Delta table
sales_table = delta_manager.create_delta_table(
    sales_df, 
    "/mnt/delta/sales", 
    partition_cols=["year", "month"]
)

# Perform daily updates
delta_manager.upsert_data(
    "/mnt/delta/sales",
    daily_updates_df,
    "target.order_id = updates.order_id"
)

# Time travel for data auditing
historical_data = delta_manager.time_travel_query(
    "/mnt/delta/sales", 
    timestamp="2025-01-01"
)

# MLOps pipeline with MLflow Model Registry
import mlflow
from mlflow.tracking import MlflowClient
from databricks.feature_store import FeatureStoreClient

class MLOpsManager:
    def __init__(self):
        self.client = MlflowClient()
        self.fs_client = FeatureStoreClient()
    
    def create_feature_store_table(self, features_df, table_name, primary_keys):
        """Create feature store table"""
        self.fs_client.create_table(
            name=table_name,
            primary_keys=primary_keys,
            df=features_df,
            description="Customer behavior features"
        )
    
    def register_model(self, model_uri, model_name):
        """Register model in MLflow registry"""
        model_version = mlflow.register_model(model_uri, model_name)
        return model_version
    
    def promote_model_stage(self, model_name, version, stage):
        """Promote model to production"""
        self.client.transition_model_version_stage(
            name=model_name,
            version=version,
            stage=stage,
            archive_existing_versions=True
        )
    
    def create_model_webhook(self, model_name, events, webhook_url):
        """Setup model lifecycle webhooks"""
        self.client.create_webhook(
            events=events,
            model_name=model_name,
            http_url_spec=mlflow.utils.rest_utils.http_request_safe({
                "url": webhook_url
            })
        )
    
    def batch_inference(self, model_name, stage, input_table):
        """Run batch inference with registered model"""
        model_uri = f"models:/{model_name}/{stage}"
        model = mlflow.spark.load_model(model_uri)
        
        # Load features
        features_df = spark.table(input_table)
        
        # Generate predictions
        predictions_df = model.transform(features_df)
        
        # Save predictions
        predictions_df.write.format("delta").mode("overwrite").saveAsTable("predictions.customer_scores")
        
        return predictions_df
    
    def model_monitoring(self, predictions_df, ground_truth_df):
        """Monitor model performance"""
        from pyspark.ml.evaluation import RegressionEvaluator
        
        # Join predictions with ground truth
        monitoring_df = predictions_df.join(ground_truth_df, "customer_id")
        
        # Calculate metrics
        evaluator = RegressionEvaluator(
            labelCol="actual_revenue",
            predictionCol="prediction",
            metricName="rmse"
        )
        
        current_rmse = evaluator.evaluate(monitoring_df)
        
        # Log monitoring metrics
        with mlflow.start_run():
            mlflow.log_metric("production_rmse", current_rmse)
            mlflow.log_metric("prediction_count", monitoring_df.count())
        
        return current_rmse

# MLOps workflow
mlops = MLOpsManager()

# Create feature store
mlops.create_feature_store_table(
    features_df, 
    "ml.customer_features", 
    ["customer_id", "date"]
)

# Register model
model_version = mlops.register_model(
    "runs:/abc123/model", 
    "customer_revenue_predictor"
)

# Promote to production
mlops.promote_model_stage(
    "customer_revenue_predictor", 
    model_version.version, 
    "Production"
)

# Setup monitoring
predictions = mlops.batch_inference(
    "customer_revenue_predictor", 
    "Production", 
    "features.customer_monthly"
)