src/customer_intelligence/api/routers/predictions.py

"""ML prediction endpoints."""

import logging

import pandas as pd
from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy import bindparam, text
from sqlalchemy.orm import Session

logger = logging.getLogger(__name__)

from customer_intelligence.api.dependencies import get_churn_model, get_db
from customer_intelligence.ml.churn import CHURN_PROBA_CUTOFF
from customer_intelligence.ml.churn import FEATURE_COLS as CHURN_FEATURES
from customer_intelligence.ml.schemas import (
    BatchChurnRequest,
    BatchChurnResponse,
    ChurnPredictRequest,
    ChurnPredictResponse,
    SegmentPredictRequest,
    SegmentPredictResponse,
    TopFactor,
)
from customer_intelligence.ml.segmentation import SEGMENT_NAMES

router = APIRouter(prefix="/predict", tags=["predictions"])

SEGMENT_FEATURE_COLS = [
    "recency_days", "frequency", "monetary",
    "avg_order_value", "tenure_days", "avg_days_between_orders",
]


def _fetch_churn_features(customer_unique_id: str, db: Session) -> dict | None:
    """Fetch pre-aggregated features for churn prediction from the warehouse."""
    row = db.execute(
        text("""
            SELECT
                dc.customer_unique_id,
                CAST(julianday(:today) - julianday(MAX(f.order_date_id)) AS INTEGER) AS recency_days,
                COUNT(DISTINCT f.order_id)           AS frequency,
                SUM(f.price + f.freight_value)       AS monetary,
                AVG(f.price + f.freight_value)       AS avg_order_value,
                AVG(f.review_score)                  AS avg_review_score,
                AVG(f.days_to_delivery)              AS avg_days_to_delivery,
                AVG(CASE WHEN f.is_late THEN 1.0 ELSE 0.0 END) AS late_delivery_rate,
                AVG(f.delivery_delay_days)           AS avg_delivery_delay,
                AVG(CASE WHEN f.review_score <= 2 THEN 1.0 ELSE 0.0 END) AS complaint_rate,
                CAST(julianday(MAX(f.order_date_id)) - julianday(MIN(f.order_date_id)) AS INTEGER) AS tenure_days
            FROM fact_orders f
            JOIN dim_customers dc ON f.customer_key = dc.customer_key
            WHERE dc.customer_unique_id = :uid
              AND f.order_status NOT IN ('canceled', 'unavailable')
            GROUP BY dc.customer_unique_id
        """),
        {"uid": customer_unique_id, "today": str(pd.Timestamp.today().date())},
    ).fetchone()
    return dict(row._mapping) if row else None


def _fetch_churn_features_batch(customer_ids: list[str], db: Session) -> dict[str, dict]:
    """Fetch churn features for multiple customers in a single query."""
    rows = db.execute(
        text("""
            SELECT
                dc.customer_unique_id,
                CAST(julianday(:today) - julianday(MAX(f.order_date_id)) AS INTEGER) AS recency_days,
                COUNT(DISTINCT f.order_id)           AS frequency,
                SUM(f.price + f.freight_value)       AS monetary,
                AVG(f.price + f.freight_value)       AS avg_order_value,
                AVG(f.review_score)                  AS avg_review_score,
                AVG(f.days_to_delivery)              AS avg_days_to_delivery,
                AVG(CASE WHEN f.is_late THEN 1.0 ELSE 0.0 END) AS late_delivery_rate,
                AVG(f.delivery_delay_days)           AS avg_delivery_delay,
                AVG(CASE WHEN f.review_score <= 2 THEN 1.0 ELSE 0.0 END) AS complaint_rate,
                CAST(julianday(MAX(f.order_date_id)) - julianday(MIN(f.order_date_id)) AS INTEGER) AS tenure_days
            FROM fact_orders f
            JOIN dim_customers dc ON f.customer_key = dc.customer_key
            WHERE dc.customer_unique_id IN :ids
              AND f.order_status NOT IN ('canceled', 'unavailable')
            GROUP BY dc.customer_unique_id
        """).bindparams(bindparam("ids", expanding=True)),
        {"ids": customer_ids, "today": str(pd.Timestamp.today().date())},
    ).fetchall()
    return {row.customer_unique_id: dict(row._mapping) for row in rows}


@router.post("/churn", response_model=ChurnPredictResponse)
def predict_churn(
    request: ChurnPredictRequest,
    db: Session = Depends(get_db),
    model=Depends(get_churn_model),
):
    features = _fetch_churn_features(request.customer_unique_id, db)
    if not features:
        raise HTTPException(status_code=404, detail="Customer not found or has no orders")

    X = pd.DataFrame([{col: float(features.get(col) or 0) for col in CHURN_FEATURES}])
    proba = float(model.predict_proba(X)[0, 1])

    # Compute SHAP feature contributions
    try:
        import shap
        explainer = shap.TreeExplainer(model)
        sv = explainer.shap_values(X.values)  # shape (1, n_features)
        top_factors = sorted(
            [TopFactor(feature=col, impact=round(float(sv[0][i]), 4)) for i, col in enumerate(CHURN_FEATURES)],
            key=lambda f: abs(f.impact),
            reverse=True,
        )[:5]
    except Exception as e:
        logger.warning("SHAP computation failed: %s", e)
        top_factors = None

    return ChurnPredictResponse(
        customer_unique_id=request.customer_unique_id,
        churn_probability=proba,
        churn_label=proba >= CHURN_PROBA_CUTOFF,
        model_version="1.0",
        top_factors=top_factors,
    )


@router.post("/churn/batch", response_model=BatchChurnResponse)
def predict_churn_batch(
    request: BatchChurnRequest,
    db: Session = Depends(get_db),
    model=Depends(get_churn_model),
):
    features_by_uid = _fetch_churn_features_batch(list(request.customer_ids), db)

    rows = [
        {col: float(features_by_uid[uid].get(col) or 0) for col in CHURN_FEATURES}
        for uid in request.customer_ids
        if uid in features_by_uid
    ]
    found_ids = [uid for uid in request.customer_ids if uid in features_by_uid]

    if not rows:
        return BatchChurnResponse(predictions=[])

    X_batch = pd.DataFrame(rows)
    probas = model.predict_proba(X_batch)[:, 1]

    predictions = [
        ChurnPredictResponse(
            customer_unique_id=uid,
            churn_probability=float(proba),
            churn_label=float(proba) >= CHURN_PROBA_CUTOFF,
            model_version="1.0",
        )
        for uid, proba in zip(found_ids, probas)
    ]
    return BatchChurnResponse(predictions=predictions)


@router.post("/segment", response_model=SegmentPredictResponse)
def predict_segment(
    request: SegmentPredictRequest,
    db: Session = Depends(get_db),
):
    # Use pre-computed segment from warehouse (fast path)
    row = db.execute(
        text("""
            SELECT customer_unique_id, segment_label
            FROM dim_customers
            WHERE customer_unique_id = :uid
            LIMIT 1
        """),
        {"uid": request.customer_unique_id},
    ).fetchone()

    if not row:
        raise HTTPException(status_code=404, detail="Customer not found")

    segment_id = row.segment_label if row.segment_label is not None else -1
    return SegmentPredictResponse(
        customer_unique_id=request.customer_unique_id,
        segment_id=segment_id,
        segment_name=SEGMENT_NAMES.get(segment_id, "Unknown"),
        model_version="precomputed",
    )


@router.get("/models/info")
def models_info():
    return {
        "churn_classifier": {"version": "1.0", "stage": "Production", "source": "disk"},
        "customer_segmentation": {"version": "1.0", "stage": "Production", "source": "disk"},
    }