"""Re-save models with importable AutoencoderWrapper."""
import os, sys
sys.path.insert(0, '/app/fraud_detection')
import joblib
import torch
from ae_model import Autoencoder, AutoencoderWrapper
from config import MODELS_DIR, DATA_DIR

# Load non-AE models
models = joblib.load(os.path.join(MODELS_DIR, "all_models.joblib"))
print(f"Loaded models: {list(models.keys())}")

# Rebuild autoencoder
data = joblib.load(os.path.join(DATA_DIR, "processed_data.joblib"))
input_dim = data['X_train'].shape[1]
ae = Autoencoder(input_dim)
ae.load_state_dict(torch.load(os.path.join(MODELS_DIR, "autoencoder.pt"), weights_only=True))
ae.eval()

models['Autoencoder'] = AutoencoderWrapper(ae)

# Save
joblib.dump(models, os.path.join(MODELS_DIR, "all_models_with_ae.joblib"))
print(f"Saved {len(models)} models with AE wrapper")
print(f"Models: {list(models.keys())}")