{#- Context variables:
   - catalog: str (e.g., "acme_fsi_demo")
   - demo_title: str (e.g., "Acme Bank FSI Demo")
-#}
# Databricks notebook source

# COMMAND ----------

# MAGIC %md
# MAGIC # {{ demo_title }} — Sample ML Training
# MAGIC
# MAGIC Trains a simple model on the sample customer data to verify MLflow + Unity Catalog registration works.
# MAGIC This file is replaced by `/plan` + `/build-next` with real industry-specific models.

# COMMAND ----------

dbutils.widgets.text("catalog", "{{ catalog }}")
catalog = dbutils.widgets.get("catalog")

# COMMAND ----------

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np

# COMMAND ----------

# MAGIC %md ## 1. Load Sample Data

# COMMAND ----------

try:
    df = spark.table(f"{catalog}.default.gold_customer_summary")
    print(f"Gold table found: {df.count()} rows")
except Exception:
    print("Gold table not ready yet — using synthetic data for model training demo")
    df = None

# Use the silver table for features if available, else generate synthetic
try:
    features_df = spark.table(f"{catalog}.default.silver_customers").toPandas()
    print(f"Loaded {len(features_df)} rows from silver_customers")
except Exception:
    print("Silver table not ready — generating synthetic features")
    np.random.seed(42)
    N = 100
    features_df = pd.DataFrame({
        "age": np.random.randint(20, 70, N),
        "balance": np.random.uniform(100, 50000, N).round(2),
        "risk_score": np.random.uniform(0.1, 0.9, N).round(3),
    })
    features_df["high_value"] = (features_df["balance"] > 25000).astype(int)

# COMMAND ----------

# MAGIC %md ## 2. Train Model

# COMMAND ----------

# Prepare features
if "high_value" not in features_df.columns:
    features_df["high_value"] = (features_df["balance"] > 25000).astype(int)

feature_cols = ["age", "balance", "risk_score"]
# Ensure numeric types
for col in feature_cols:
    if col in features_df.columns:
        features_df[col] = pd.to_numeric(features_df[col], errors="coerce")

features_df = features_df.dropna(subset=feature_cols + ["high_value"])
X = features_df[feature_cols]
y = features_df["high_value"]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

mlflow.sklearn.autolog()

with mlflow.start_run(run_name="sample_model_training"):
    model = RandomForestClassifier(n_estimators=10, max_depth=3, random_state=42)
    model.fit(X_train, y_train)
    accuracy = model.score(X_test, y_test)
    print(f"Model accuracy: {accuracy:.3f}")

print("Sample model training complete — MLflow run logged.")

# COMMAND ----------

# MAGIC %md ## 3. Register Model (optional — uncomment after first successful run)

# COMMAND ----------

# Uncomment to register in Unity Catalog:
# model_uri = f"runs:/{mlflow.last_active_run().info.run_id}/model"
# mlflow.register_model(model_uri, f"{catalog}.default.sample_classifier")
# print("Model registered in Unity Catalog")
