import joblib
import pandas as pd
# Импортируем наши классы, чтобы joblib мог их десериализовать
from train import TextExtractor, NumberExtractor

# 1. Загружаем модель
model = joblib.load('solution/model/mcc_model.pkl')

# 2. Подготавливаем тестовые данные (как они придут в API)
test_json = {
    "transaction_id": "TX00001116",
    "terminal_name": "STORE001",
    "terminal_description": "common common common thing",
    "city": "NYC",
    "amount": 272.80,
    "items": [
        {"name": "basic loyalty", "price": 58.20},
        {"name": "Bringiong item lifes", "price": 28.99}
    ]
}


# 3. Преобразуем в формат, который понимает пайплайн (как в main.py)
items_str = " ".join([i['name'] for i in test_json['items']])
full_text = f"{test_json['terminal_name']} {test_json['terminal_description']} {items_str}".lower()

input_df = pd.DataFrame([{
    'full_text': full_text,
    'amount': test_json['amount']
}])

# 4. Делаем предсказание
prediction = model.predict(input_df)[0]
confidence = model.predict_proba(input_df).max()

print(f"Transaction ID: {test_json['transaction_id']}")
print(f"Predicted MCC: {prediction}")
print(f"Confidence: {confidence:.4f}")