"""
Model training module.
Trains a LightGBM classifier with cross-validation, feature selection,
and regularization for balanced accuracy. Prototype-grade.
"""

import joblib
import numpy as np
import pandas as pd
from sklearn.model_selection import TimeSeriesSplit, cross_val_score
from sklearn.metrics import (
    roc_auc_score,
    precision_score,
    recall_score,
    f1_score,
    confusion_matrix,
    classification_report,
    mean_absolute_error,
    r2_score,
)
from lightgbm import LGBMClassifier, LGBMRegressor

from salesforce_client import connect_to_salesforce, fetch_promotion_data
from data_preprocessing import preprocess, get_feature_columns, compute_roi_label, get_roi_feature_columns

MODEL_PATH            = "model.joblib"
ROI_MODEL_PATH        = "roi_model.joblib"
SERVING_CONTEXT_PATH  = "serving_context.joblib"


def select_features(model, feature_cols, X_train, y_train, min_importance_pct=0.5):
    """
    Drop features with near-zero importance to reduce noise and overfitting.
    Keeps features whose importance is above min_importance_pct of total.
    """
    importances = model.feature_importances_
    total_importance = importances.sum()

    if total_importance == 0:
        return feature_cols, list(range(len(feature_cols)))

    importance_pct = (importances / total_importance) * 100
    selected_indices = [i for i, pct in enumerate(importance_pct) if pct >= min_importance_pct]

    if len(selected_indices) < 5:
        # Keep at least top 10 features
        selected_indices = list(np.argsort(importances)[::-1][:10])

    selected_features = [feature_cols[i] for i in selected_indices]

    dropped = len(feature_cols) - len(selected_features)
    if dropped > 0:
        print(f"  Feature selection: kept {len(selected_features)}/{len(feature_cols)} "
              f"(dropped {dropped} low-importance features)")

    return selected_features, selected_indices


def train_and_evaluate():
    """End-to-end pipeline: fetch, enrich, preprocess, train, validate, save."""

    # Step 1: Connect and fetch data
    print("=" * 60)
    print("STEP 1: Fetching data from Salesforce")
    print("=" * 60)
    sf = connect_to_salesforce()
    df = fetch_promotion_data(sf)

    if df.empty:
        print("No data returned from Salesforce. Exiting.")
        return

    # Step 2: Preprocess (labels, missing values, feature engineering, encoding)
    print("\n" + "=" * 60)
    print("STEP 2: Preprocessing data")
    print("=" * 60)
    df, encoders = preprocess(df)

    # Step 3: Temporal sort — required for leakage-safe holdout (PROJECT_CONTEXT §7).
    # pds_created_date = PDS Header CreatedDate, set in salesforce_client._flatten_record.
    df["_sort_date"] = pd.to_datetime(df.get("pds_created_date"), errors="coerce")
    n_nat = df["_sort_date"].isna().sum()
    if n_nat > 0:
        print(f"  WARNING: {n_nat} records have no date and are excluded from training.")
        df = df[df["_sort_date"].notna()].copy()
    if df.empty:
        print("No records with valid dates. Exiting.")
        return
    df = df.sort_values("_sort_date").reset_index(drop=True)

    # Step 4: Prepare features and labels
    feature_cols = get_feature_columns()
    for col in feature_cols:
        if col not in df.columns:
            df[col] = 0.0

    X = df[feature_cols]
    y = df["label"]

    print(f"\nDataset shape: {X.shape}")
    print(f"Total features: {len(feature_cols)}")
    print(f"Positive class (success): {y.sum()} ({y.mean():.1%})")
    print(f"Negative class (failure): {(1 - y).sum()} ({(1 - y).mean():.1%})")

    # Step 5: Temporal train/test split (earliest 80% -> train, most recent 20% -> test).
    # Temporal split is required so the model is evaluated on data it never saw during
    # training, matching real deployment conditions. Random stratified splits inflate AUC.
    split_idx = int(len(df) * 0.8)
    X_train, X_test = X.iloc[:split_idx], X.iloc[split_idx:]
    y_train, y_test = y.iloc[:split_idx], y.iloc[split_idx:]
    df_train = df.iloc[:split_idx]

    train_date_min = df["_sort_date"].iloc[0].date()
    train_date_max = df["_sort_date"].iloc[split_idx - 1].date()
    test_date_min  = df["_sort_date"].iloc[split_idx].date()
    test_date_max  = df["_sort_date"].iloc[-1].date()
    print(f"Train: {len(X_train)} records ({train_date_min} -> {train_date_max})")
    print(f"Test:  {len(X_test)}  records ({test_date_min} -> {test_date_max})")
    print(f"Train positive rate: {y_train.mean():.1%} | Test positive rate: {y_test.mean():.1%}")

    if y_test.nunique() < 2:
        print("WARNING: test split contains only one class — ROC-AUC is undefined. "
              "Consider using a 70/30 split or collecting more recent failure examples.")
        return

    # Step 6: Train initial model with regularization
    print("\n" + "=" * 60)
    print("STEP 4: Training LightGBM (with regularization)")
    print("=" * 60)

    n_negative = (y_train == 0).sum()
    n_positive = (y_train == 1).sum()
    scale_pos_weight = n_negative / max(n_positive, 1)

    model = LGBMClassifier(
        n_estimators=400,
        learning_rate=0.03,
        max_depth=6,
        num_leaves=31,
        min_child_samples=20,
        # Regularization to prevent overfitting with more features
        reg_alpha=0.1,        # L1 regularization
        reg_lambda=1.0,       # L2 regularization
        colsample_bytree=0.8, # Use 80% of features per tree
        subsample=0.8,        # Use 80% of data per tree
        subsample_freq=5,
        scale_pos_weight=scale_pos_weight,
        random_state=42,
        verbose=-1,
    )
    model.fit(X_train, y_train)
    print("Initial model training complete.")

    # Step 7: Feature selection - drop noise features
    print("\n" + "=" * 60)
    print("STEP 5: Feature selection")
    print("=" * 60)

    selected_features, selected_indices = select_features(
        model, feature_cols, X_train, y_train, min_importance_pct=0.3
    )

    # Retrain with selected features if any were dropped
    if len(selected_features) < len(feature_cols):
        X_train_sel = X_train[selected_features]
        X_test_sel = X_test[selected_features]

        model_sel = LGBMClassifier(
            n_estimators=400,
            learning_rate=0.03,
            max_depth=6,
            num_leaves=31,
            min_child_samples=20,
            reg_alpha=0.1,
            reg_lambda=1.0,
            colsample_bytree=0.8,
            subsample=0.8,
            subsample_freq=5,
            scale_pos_weight=scale_pos_weight,
            random_state=42,
            verbose=-1,
        )
        model_sel.fit(X_train_sel, y_train)

        # Compare: only keep selected if it doesn't hurt performance
        auc_full = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])
        auc_sel = roc_auc_score(y_test, model_sel.predict_proba(X_test_sel)[:, 1])

        print(f"  Full model AUC:     {auc_full:.4f}")
        print(f"  Selected model AUC: {auc_sel:.4f}")

        if auc_sel >= auc_full - 0.005:  # Allow 0.5% tolerance
            model = model_sel
            feature_cols = selected_features
            X_train = X_train_sel
            X_test = X_test_sel
            print(f"  -> Using selected features ({len(selected_features)} features)")
        else:
            print(f"  -> Keeping all features (selection hurt AUC by {auc_full - auc_sel:.4f})")

    # Step 8: Walk-forward cross-validation within the training slice.
    # TimeSeriesSplit respects temporal order: each fold trains on older data and
    # validates on newer data, matching the holdout philosophy.
    print("\n" + "=" * 60)
    print("STEP 6: Walk-forward cross-validation (5-fold, TimeSeriesSplit)")
    print("=" * 60)

    cv = TimeSeriesSplit(n_splits=5)
    cv_scores = cross_val_score(model, X_train, y_train, cv=cv, scoring="roc_auc")
    print(f"  CV AUC scores: {[f'{s:.4f}' for s in cv_scores]}")
    print(f"  CV AUC mean:   {cv_scores.mean():.4f} (+/- {cv_scores.std():.4f})")

    # Check for overfitting: train AUC vs CV AUC
    train_auc = roc_auc_score(y_train, model.predict_proba(X_train)[:, 1])
    print(f"  Train AUC:     {train_auc:.4f}")
    overfit_gap = train_auc - cv_scores.mean()
    if overfit_gap > 0.05:
        print(f"  WARNING: Potential overfitting (gap={overfit_gap:.4f}). "
              f"Consider stronger regularization.")
    else:
        print(f"  Overfit gap:   {overfit_gap:.4f} (healthy)")

    # Step 9: Final evaluation on held-out test set
    print("\n" + "=" * 60)
    print("STEP 7: Test set evaluation")
    print("=" * 60)

    y_pred = model.predict(X_test)
    y_proba = model.predict_proba(X_test)[:, 1]

    auc = roc_auc_score(y_test, y_proba)
    precision = precision_score(y_test, y_pred)
    recall = recall_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred)
    cm = confusion_matrix(y_test, y_pred)

    print(f"\nROC-AUC:   {auc:.4f}")
    print(f"Precision: {precision:.4f}")
    print(f"Recall:    {recall:.4f}")
    print(f"F1-Score:  {f1:.4f}")
    print(f"\nConfusion Matrix:")
    print(f"  TN={cm[0][0]}  FP={cm[0][1]}")
    print(f"  FN={cm[1][0]}  TP={cm[1][1]}")
    print(f"\nClassification Report:")
    print(classification_report(y_test, y_pred, target_names=["Failure", "Success"]))

    # Step 10: Feature importance analysis
    print("Feature Importances (all):")
    importances = model.feature_importances_
    total = importances.sum()
    sorted_idx = np.argsort(importances)[::-1]
    for rank, i in enumerate(sorted_idx):
        pct = (importances[i] / total) * 100 if total > 0 else 0
        print(f"  {rank+1:2d}. {feature_cols[i]:45s} {importances[i]:6.0f} ({pct:5.1f}%)")
        if rank >= 19:
            remaining = len(sorted_idx) - 20
            if remaining > 0:
                print(f"  ... and {remaining} more features")
            break

    # Step 11: Save acceptance model
    print("\n" + "=" * 60)
    print("STEP 8: Saving model")
    print("=" * 60)
    joblib.dump(
        {"model": model, "encoders": encoders, "features": feature_cols},
        MODEL_PATH,
    )
    print(f"Model saved to {MODEL_PATH}")
    print(f"Features: {len(feature_cols)} | Estimators: {model.n_estimators}")

    # Step 12: Train ROI model
    print("\n" + "=" * 60)
    print("STEP 9: Training ROI model (quantity achievement regressor)")
    print("=" * 60)
    train_roi_model(df, feature_cols, encoders)

    # Step 13: Serialize serving context (benchmark stats + analogues index).
    # Benchmark stats are computed on the training slice only so they reflect the
    # same population the model learned from. Analogues index uses the full dataset
    # so the similarity search covers all historical records.
    print("\n" + "=" * 60)
    print("STEP 10: Building and saving serving_context.joblib")
    print("=" * 60)
    _save_serving_context(df_train, analogues_df=df)


def train_roi_model(df: pd.DataFrame, feature_cols: list, encoders: dict):
    """
    Two-stage ROI model:
      Stage 1 — acceptance model (already trained, AUC ~0.987).
      Stage 2 — this regressor: predicts achievement ratio CONDITIONED ON acceptance.
                Trained only on accepted promotions so it learns the variance in
                how well accepted promotions actually perform, without the distortion
                of rejected promotions (ratio=0) pulling the model toward predicting
                low values for everything.

    At inference:  expected_units = acceptance_probability * achievement_ratio * qty_target
    Label source:  SellInFact / SellInTarget (100% filled, perfect corr with QuantityFact/Target)
    """
    from salesforce_client import SUCCESS_STATUSES

    # cpd_sell_in_fact/target are flattened from CPD relationship by salesforce_client.
    # They are only non-null for Validated PDSD lines after the Synchro batch runs.
    if ("cpd_sell_in_fact" not in df.columns
            or df["cpd_sell_in_fact"].isna().all()):
        print("  cpd_sell_in_fact not available — CPD relationship fields absent from query.")
        print("  ROI predictions at inference time will use the heuristic formula.")
        return

    # ── Stage 2: train on ACCEPTED promotions only ────────────────────────────
    accepted_mask = (
        df["Negoptim__Status__c"].isin(SUCCESS_STATUSES)
        & df["cpd_sell_in_fact"].notna()
        & df["cpd_sell_in_target"].notna()
        & (df["cpd_sell_in_target"] > 0)
    )

    # Supplementary fields only included when present (CPD-era fields, absent on PDSD)
    NEW_ROI_FIELDS = [
        "Negoptim__SellInTarget__c",
        "Negoptim__SellOutDiscMassValueContribTarget__c",
        "Negoptim__SellInDiscountMassValueTarget__c",
        "Negoptim__SalesExecutionRateTarget__c",
        "Negoptim__TariffCostRegular__c",
        "Negoptim__SellInBDateMonth__c",
        "Negoptim__SellInBDateWeek__c",
    ]
    roi_df = df[accepted_mask].copy()
    roi_df["qty_achievement_ratio"] = (
        roi_df["cpd_sell_in_fact"] / roi_df["cpd_sell_in_target"]
    ).clip(0, 5)

    print(f"  Stage 2 training records (accepted only): {len(roi_df)}")

    if len(roi_df) < 50:
        print("  Not enough ROI records — skipping ROI model training.")
        return

    # Remove extreme outliers (ratio > 3 — likely data entry errors)
    roi_df = roi_df[roi_df["qty_achievement_ratio"] <= 3.0]
    print(f"  After outlier removal (ratio <= 3): {len(roi_df)}")

    y_raw = roi_df["qty_achievement_ratio"]
    print(f"  Achievement ratio: mean={y_raw.mean():.2f} | median={y_raw.median():.2f} "
          f"| std={y_raw.std():.2f}")

    # Build feature set: base ROI features + new Salesforce fields
    base_cols = get_roi_feature_columns()
    roi_feature_cols = base_cols + [
        f for f in NEW_ROI_FIELDS if f in roi_df.columns and f not in base_cols
    ]
    roi_feature_cols = [c for c in roi_feature_cols if c in roi_df.columns]

    X_roi = roi_df[roi_feature_cols].fillna(0.0)

    # Log-transform target to handle right skew
    y_log = np.log1p(y_raw)

    # Temporal split — sort roi_df by pds_created_date (already sorted in the
    # parent df; accepted_mask preserves that order).
    roi_split = int(len(roi_df) * 0.8)
    X_train_r  = X_roi.iloc[:roi_split]
    X_test_r   = X_roi.iloc[roi_split:]
    y_train_r  = y_raw.iloc[:roi_split]
    y_test_r   = y_raw.iloc[roi_split:]
    y_log_train = y_log.iloc[:roi_split]

    # Train three quantile regressors (p10, p50, p90) so the API can return
    # a profit range instead of a point estimate (PROJECT_CONTEXT §2 condition 2).
    quantile_models = {}
    _base_params = dict(
        n_estimators=500,
        learning_rate=0.03,
        max_depth=6,
        num_leaves=31,
        min_child_samples=15,
        reg_alpha=0.1,
        reg_lambda=1.0,
        colsample_bytree=0.8,
        subsample=0.8,
        random_state=42,
        verbose=-1,
    )
    for alpha, key in [(0.1, "p10"), (0.5, "p50"), (0.9, "p90")]:
        m = LGBMRegressor(objective="quantile", alpha=alpha, **_base_params)
        m.fit(X_train_r, y_log_train)
        quantile_models[key] = m

    # Report evaluation metrics on the median model
    y_pred_log = quantile_models["p50"].predict(X_test_r)
    y_pred = np.expm1(y_pred_log)
    mae = mean_absolute_error(y_test_r, y_pred)
    r2  = r2_score(y_test_r, y_pred)

    print(f"  MAE (achievement ratio, p50): {mae:.3f}")
    print(f"  R2  (conditioned on acceptance, p50): {r2:.3f}")
    print(f"  Interpretation: on average, prediction is off by {mae*100:.1f}% of target qty")

    # Feature importance from the p50 model
    importances = quantile_models["p50"].feature_importances_
    total = importances.sum()
    sorted_idx = np.argsort(importances)[::-1]
    print("  Top 10 features for ROI model (p50):")
    for rank, i in enumerate(sorted_idx[:10]):
        pct = (importances[i] / total) * 100 if total > 0 else 0
        print(f"    {rank+1:2d}. {roi_feature_cols[i]:45s} {pct:5.1f}%")

    joblib.dump(
        {
            "models": quantile_models,
            "features": roi_feature_cols,
            "log_transform": True,
            "conditioned_on_acceptance": True,
        },
        ROI_MODEL_PATH,
    )
    print(f"  ROI model saved to {ROI_MODEL_PATH} (p10/p50/p90 quantile models)")


def _save_serving_context(df: pd.DataFrame, analogues_df: pd.DataFrame = None) -> None:
    """
    Build and serialize all Salesforce-derived state the API needs at startup.
    Lets the HF Space serve predictions without a live Salesforce connection.

    df           — training slice only; benchmark stats and enrichment context are
                   computed here so they reflect the same population the model trained on.
    analogues_df — full dataset (train + test); analogue index uses all records so
                   similarity search has the widest possible coverage.
    """
    from salesforce_client import SUCCESS_STATUSES
    from data_preprocessing import handle_missing_values, engineer_features

    if analogues_df is None:
        analogues_df = df

    # ── 1. Benchmark stats (training slice only) ──────────────────────────────
    labels = df["Negoptim__Status__c"].apply(lambda s: 1 if s in SUCCESS_STATUSES else 0)
    benchmark_stats = {
        "average_success_rate": round(float(labels.mean()), 4),
        "total_records":        int(len(df)),
        "success_count":        int(labels.sum()),
        "failure_count":        int((1 - labels).sum()),
    }
    print(f"  Benchmark (training slice): {benchmark_stats['total_records']} records, "
          f"{benchmark_stats['average_success_rate']:.1%} success rate")

    # ── 2. Enrichment context (training slice only) ───────────────────────────
    df_enc = df.copy()
    df_enc["_label"] = labels
    enrichment_context = _build_enrichment_context_for_serving(df_enc)
    print(f"  Enrichment: {len(enrichment_context.get('pg_stats', {}))} groups")

    # ── 3. Analogues index (full dataset) ────────────────────────────────────
    ANALOGUE_FEATURES = [
        "Negoptim__SellOutDiscountPerc__c",
        "Negoptim__SellOutDiscountUAmt__c",
        "Negoptim__SellOutDiscountQtyTarget__c",
        "Negoptim__SellOutSourceContribPerc__c",
        "Negoptim__SellOutExecutionRateTarget__c",
        "Negoptim__Perc__c",
        "discount_depth_vs_gross",
        "campaign_duration_days",
        "sell_out_month",
        "sell_out_quarter",
    ]

    adf = analogues_df.copy()
    adf = handle_missing_values(adf)
    adf = engineer_features(adf)

    present = [f for f in ANALOGUE_FEATURES if f in adf.columns]
    mat = adf[present].fillna(0.0).values.astype(np.float64)
    norms = np.linalg.norm(mat, axis=1, keepdims=True)
    norms = np.where(norms < 1e-9, 1.0, norms)
    analogues_matrix = mat / norms

    sell_out_dates = pd.to_datetime(adf.get("Negoptim__SellOutBDate__c"), errors="coerce")
    analogues_meta = pd.DataFrame({
        "status":                 adf["Negoptim__Status__c"].apply(
                                      lambda s: "SUCCESS" if s in SUCCESS_STATUSES else "FAILURE"),
        "retailer":               adf.get("retailer_name", pd.Series(["UNKNOWN"] * len(adf))).fillna("UNKNOWN").astype(str),
        "supplier":               adf.get("Negoptim__Supplier__c", pd.Series(["Unknown"] * len(adf))).fillna("Unknown").astype(str),
        "discount_pct":           pd.to_numeric(adf.get("Negoptim__SellOutDiscountPerc__c", 0), errors="coerce").fillna(0),
        "qty_target":             pd.to_numeric(adf.get("Negoptim__SellOutDiscountQtyTarget__c", 0), errors="coerce").fillna(0),
        "campaign_duration_days": pd.to_numeric(adf.get("campaign_duration_days", 0), errors="coerce").fillna(0),
        "discount_depth_vs_gross":pd.to_numeric(adf.get("discount_depth_vs_gross", 0), errors="coerce").fillna(0).round(4),
        "supplier_acceptance_rate":pd.to_numeric(adf.get("supplier_acceptance_rate", 0), errors="coerce").fillna(0).round(4),
        "sell_out_date":          sell_out_dates.dt.strftime("%Y-%m-%d").fillna(""),
        "n_features":             len(present),
    })
    print(f"  Analogues: {len(analogues_meta)} records × {len(present)} features")

    # ── 4. Save ───────────────────────────────────────────────────────────────
    joblib.dump(
        {
            "benchmark_stats":    benchmark_stats,
            "enrichment_context": enrichment_context,
            "analogues_matrix":   analogues_matrix,
            "analogues_meta":     analogues_meta,
            "analogue_features":  present,
        },
        SERVING_CONTEXT_PATH,
    )
    print(f"  serving_context.joblib saved ({len(present)} analogue features)")


def _build_enrichment_context_for_serving(df: pd.DataFrame) -> dict:
    """Replica of api._build_enrichment_context — avoids circular imports."""
    from salesforce_client import SUCCESS_STATUSES

    context = {"pg_stats": {}, "monthly_demand": {}, "bu_id_map": {}, "achievement": {}}

    # bu_id_map: PDSD uses BusinessUnit__c (Salesforce ID); no retailer_name available
    if "Negoptim__BusinessUnit__c" in df.columns:
        bu_ids = df["Negoptim__BusinessUnit__c"].dropna().unique()
        context["bu_id_map"] = {str(b): str(b) for b in bu_ids}

    df = df.copy()
    df["_date"] = pd.to_datetime(df.get("Negoptim__SellOutBDate__c"), errors="coerce")
    # Discount depth from SellOutDiscountPerc__c (field is 0–100 scale)
    disc_pct = pd.to_numeric(df.get("Negoptim__SellOutDiscountPerc__c", 0), errors="coerce").fillna(0)
    df["_discount"] = (disc_pct / 100.0).clip(0, 1)

    pg_col = "Negoptim__Supplier__c"
    if pg_col in df.columns:
        for pg, group in df.groupby(pg_col):
            dates     = group["_date"].dropna()
            lab       = group["_label"]
            discounts = group["_discount"]
            qtys      = pd.to_numeric(group.get("Negoptim__SellOutDiscountQtyTarget__c", 0), errors="coerce").fillna(0)

            pg_data = {
                "success_rate":  round(float(lab.mean()), 4),
                "count":         len(group),
                "avg_qty":       round(float(qtys.mean()), 2),
                "avg_discount":  round(float(discounts.mean()), 4),
                "recent_dates":  sorted(dates.dt.strftime("%Y-%m-%d").tolist())[-20:],
                "monthly_success": {},
            }
            if not dates.empty:
                gm = group.copy()
                gm["_month"] = gm["_date"].dt.month
                for month, mgrp in gm.groupby("_month"):
                    pg_data["monthly_success"][int(month)] = round(float(mgrp["_label"].mean()), 4)

            if len(discounts) >= 5 and discounts.std() > 0.001:
                pg_data["price_elasticity"] = round(float(np.corrcoef(discounts, lab)[0, 1]), 4)
            else:
                pg_data["price_elasticity"] = 0.0

            context["pg_stats"][str(pg)] = pg_data

    dates_all = df["_date"].dropna()
    if not dates_all.empty:
        monthly = dates_all.dt.to_period("M").value_counts().sort_index()
        avg_count = monthly.mean()
        for period, count in monthly.items():
            key = f"{period.year}-{period.month:02d}"
            context["monthly_demand"][key] = round(min((count / (avg_count + 0.01)) * 50, 100), 1)

    # Achievement context — PDSD has no QuantityFact; skip gracefully if absent
    bu_col       = "Negoptim__BusinessUnit__c"
    supplier_col = "Negoptim__Supplier__c"
    raw_fact = df.get("Negoptim__QuantityFact__c")
    if raw_fact is not None and pd.to_numeric(raw_fact, errors="coerce").notna().any():
        qty_target = pd.to_numeric(df.get("Negoptim__SellOutDiscountQtyTarget__c", 0), errors="coerce").clip(lower=1)
        qty_fact   = pd.to_numeric(raw_fact, errors="coerce")
        df["_ach"] = (qty_fact / qty_target).clip(0, 5)
        has_data = df[df["_ach"].notna()].copy()

        ach = {"retailer_supplier": {}, "retailer": {}, "supplier": {}}
        for (r, s), grp in has_data.groupby([bu_col, supplier_col]):
            vals = grp["_ach"].tolist()
            ach["retailer_supplier"][(str(r), str(s))] = {
                "avg":  round(float(np.mean(vals)), 4),
                "last": round(float(vals[-1]), 4),
                "n":    len(vals),
                "std":  round(float(np.std(vals)), 4) if len(vals) > 1 else 0.0,
            }
        for r, grp in has_data.groupby(bu_col):
            ach["retailer"][str(r)] = round(float(grp["_ach"].mean()), 4)
        for s, grp in has_data.groupby(supplier_col):
            ach["supplier"][str(s)] = round(float(grp["_ach"].mean()), 4)
        context["achievement"] = ach

    # ── Historical aggregate lookups — Fix 2 ─────────────────────────────────
    # These three dicts let api.py populate the six historical-aggregate features
    # (supplier_acceptance_rate etc.) at serving time without any Apex-side SOQL.
    # Keys that don't appear at serving fall back to dataset averages in api.py.

    dataset_avg_success = round(float(df["_label"].mean()), 4)

    # negoscope_stats: {negoscope_id: {success_rate, count}}
    scope_col = "Negoptim__NegoScope__c"
    negoscope_stats: dict = {}
    if scope_col in df.columns:
        for scope, grp in df.groupby(scope_col):
            negoscope_stats[str(scope)] = {
                "success_rate": round(float(grp["_label"].mean()), 4),
                "count":        len(grp),
            }
    context["negoscope_stats"] = negoscope_stats

    # mechanic_family_stats: {family_code: {success_rate, count}}
    # family_code = first 2 chars of SellOutDiscountType__c (e.g. "RI", "CR", "VR")
    disc_type_col = "Negoptim__SellOutDiscountType__c"
    mechanic_family_stats: dict = {}
    if disc_type_col in df.columns:
        families = df[disc_type_col].fillna("Unknown").str[:2]
        for fam, grp in df.groupby(families):
            mechanic_family_stats[str(fam)] = {
                "success_rate": round(float(grp["_label"].mean()), 4),
                "count":        len(grp),
            }
    context["mechanic_family_stats"] = mechanic_family_stats

    # supplier_product_stats: {"supplier_id|product_id": {count, last_date_iso}}
    product_col = "Negoptim__Product__c"
    supplier_product_stats: dict = {}
    if product_col in df.columns and supplier_col in df.columns:
        df["_sp_date"] = pd.to_datetime(df.get("pds_created_date"), errors="coerce")
        for (sup, prod), grp in df.groupby([supplier_col, product_col]):
            key = f"{sup}|{prod}"
            last_date = grp["_sp_date"].dropna().max()
            supplier_product_stats[key] = {
                "count":         len(grp),
                "last_date_iso": last_date.isoformat() if pd.notna(last_date) else None,
            }
    context["supplier_product_stats"] = supplier_product_stats
    context["dataset_avg_success"] = dataset_avg_success

    n_scope = len(negoscope_stats)
    n_fam   = len(mechanic_family_stats)
    n_sp    = len(supplier_product_stats)
    print(f"  Historical aggregates: {n_scope} negoscopes, {n_fam} mechanic families, "
          f"{n_sp} supplier×product pairs")

    return context


if __name__ == "__main__":
    train_and_evaluate()