"""
save_models.py  ─  TemporalDrift-ETM  |  Run this cell inside your Jupyter notebook
               after training to export all artifacts needed for Hugging Face deployment.

VARIABLE NAME MAP  (adjust the left side to match your notebook):
──────────────────────────────────────────────────────────────────
  ensemble_model  ←  your VotingClassifier / trained ensemble
  scaler          ←  StandardScaler fitted on training data
  le              ←  LabelEncoder fitted on y_train
  feature_cols    ←  list of feature column names used in training
  X_train_scaled  ←  scaled training feature matrix  (numpy array)
  y_train_encoded ←  encoded training labels          (numpy array, integers)
  y_train_labels  ←  decoded training labels          (numpy array, strings)
"""

import os
import numpy as np
import joblib

# ── OUTPUT DIRECTORY ──────────────────────────────────────────────────────────
SAVE_DIR = "models"           # relative to where you run this script
os.makedirs(SAVE_DIR, exist_ok=True)

# ── ① RENAME YOUR VARIABLES HERE ─────────────────────────────────────────────
# (Only change the right-hand side to match your notebook's variable names)

# _ensemble      = ensemble_model    # your trained VotingClassifier / ensemble
# _scaler        = scaler            # your fitted StandardScaler
# _le            = le                # your fitted LabelEncoder
# _feature_names = feature_cols      # list[str] — columns used during training
# _X_train       = X_train_scaled    # np.ndarray, shape (n_samples, n_features) — SCALED
# _y_encoded     = y_train_encoded   # np.ndarray, shape (n_samples,) — integer labels
# _y_labels      = y_train_labels    # np.ndarray, shape (n_samples,) — string labels

# ──── Replace with ────
#
# ⚠️  CRITICAL — use the BASE sklearn VotingClassifier, NOT a custom wrapper.
#
#     DO NOT use:  _ensemble = ensemble_retrained
#     WHY:  ensemble_retrained is an EnsembleWrapperRetrained object defined
#           only in this notebook's __main__ scope.  joblib pickles a *reference*
#           to that class, so loading it anywhere outside this notebook raises:
#               "Can't get attribute 'EnsembleWrapperRetrained'
#                on <module '__main__' from '/app/app.py'>"
#           AND the object serialises with no model state (55 bytes) — all
#           trained weights are silently lost.
#
#     USE:  The raw sklearn VotingClassifier (ensemble / voting_clf / etc.).
#           It contains .estimators_ (RF + XGBoost + LightGBM + CatBoost)
#           and pickles cleanly because all sklearn/xgboost classes are
#           importable from their own packages on any machine.
#
_ensemble      = ensemble              # ← the base VotingClassifier from training
_scaler        = scaler               # already correct
_le            = le                   # already correct
_feature_names = FEATURE_COLS        # capital, matches your notebook
_X_train       = X_sm_scaled[:, top_idx]   # SMOTE-resampled + feature-selected
_y_encoded     = y_train             # integer labels from Cell 16/17
_y_labels      = le.inverse_transform(y_train)

# ── ② VALIDATE before saving ──────────────────────────────────────────────────
# Catches the wrapper-class mistake immediately instead of at deploy time.
from sklearn.base import BaseEstimator

if not isinstance(_ensemble, BaseEstimator):
    raise TypeError(
        f"\n\n_ensemble is {type(_ensemble).__name__!r}, not a sklearn BaseEstimator.\n"
        "You are probably passing a custom wrapper class.  Assign the raw\n"
        "VotingClassifier (or equivalent) to _ensemble instead.\n"
        "Check: hasattr(_ensemble, 'estimators_') should be True.\n"
    )

print(f"✅  _ensemble type : {type(_ensemble).__name__}")
print(f"    estimators     : {[type(e).__name__ for e in _ensemble.estimators_]}")
print(f"    classes_       : {list(_ensemble.classes_)}")

# ── ③ SAVE CORE ARTIFACTS ─────────────────────────────────────────────────────
print("\nSaving ensemble model …")
joblib.dump(_ensemble,      os.path.join(SAVE_DIR, "ensemble_model.pkl"),  compress=3)

print("Saving scaler …")
joblib.dump(_scaler,        os.path.join(SAVE_DIR, "scaler.pkl"),          compress=3)

print("Saving label encoder …")
joblib.dump(_le,            os.path.join(SAVE_DIR, "label_encoder.pkl"),   compress=3)

print("Saving feature names …")
joblib.dump(list(_feature_names), os.path.join(SAVE_DIR, "feature_names.pkl"))

# ── ④ BUILD BASELINE PROFILES ─────────────────────────────────────────────────
# Each family profile stores:
#   • feature_samples  — per-feature value subsample (for JS-divergence drift detection)
#   • sample_X         — feature matrix subsample (for combined retraining)
#   • sample_y         — integer label subsample   (for combined retraining)
#   • n_total          — total training samples for this family

SUBSAMPLE = 500          # samples retained per family  (balance size vs accuracy)
SEED      = 42

rng = np.random.RandomState(SEED)
baseline_profiles = {}
families = np.unique(_y_labels)

print(f"\nBuilding baseline profiles for {len(families)} families …")

for family in families:
    mask = (_y_labels == family)
    fX   = _X_train[mask]
    fy   = _y_encoded[mask]
    n    = min(SUBSAMPLE, len(fX))
    idx  = rng.choice(len(fX), n, replace=False)

    baseline_profiles[family] = {
        # Dict of feature_name → list of sample values (for JS divergence)
        "feature_samples": {
            feat: fX[idx, i].tolist()
            for i, feat in enumerate(_feature_names)
        },
        "sample_X": fX[idx],      # np.ndarray (n, d) — for retraining
        "sample_y": fy[idx],      # np.ndarray (n,)   — for retraining
        "n_total":  int(mask.sum()),
    }
    print(f"  [{family}]  total={mask.sum():>6,}  saved={n}")

print("\nSaving baseline profiles …")
joblib.dump(baseline_profiles, os.path.join(SAVE_DIR, "baseline_profiles.pkl"), compress=3)

# ── ⑤ SUMMARY ─────────────────────────────────────────────────────────────────
print("\n" + "═" * 55)
print("  All artifacts saved to:", os.path.abspath(SAVE_DIR))
print("═" * 55)
for fname in sorted(os.listdir(SAVE_DIR)):
    if fname.endswith(".pkl"):
        size_mb = os.path.getsize(os.path.join(SAVE_DIR, fname)) / 1e6
        print(f"   {fname:<30}  {size_mb:>7.2f} MB")
print("═" * 55)
print("\nNext step: copy the 'models/' folder into your HF Space repo.")