import time
import joblib
import pandas as pd
import numpy as np
import random

from sklearn.base import BaseEstimator, TransformerMixin


# ===== КРИТИЧЕСКИ ВАЖНО =====
# Эти классы ДОЛЖНЫ существовать до joblib.load()

class TextExtractor(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None): 
        return self

    def transform(self, X):
        return X['full_text'].fillna('')


class NumberExtractor(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None): 
        return self

    def transform(self, X):
        return X[['amount']].fillna(0)


# ============================

MODEL_PATH = "mcc_model.pkl"

BASE_TX = {
    "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},
        {"name": "regular item basic item", "price": 56.91}
    ]
}


def json_to_df(data_list):
    rows = []
    for d in data_list:
        items = d.get("items", [])
        item_text = " ".join(str(i.get("name","")) for i in items)
        full_text = f"{d.get('terminal_name','')} {d.get('terminal_description','')} {item_text}".lower()

        rows.append({
            "full_text": full_text,
            "amount": float(d.get("amount", 0))
        })

    return pd.DataFrame(rows)


def make_batch(n=100):
    batch = []
    for i in range(n):
        tx = BASE_TX.copy()
        tx["transaction_id"] = f"TX{i:06d}"
        tx["amount"] = round(random.uniform(3, 500), 2)
        batch.append(tx)
    return batch


# ============================
#           BENCH
# ============================

print("\n[1] LOADING MODEL (cold start)...")
t0 = time.perf_counter()
model = joblib.load(MODEL_PATH)
t1 = time.perf_counter()
print(f"Cold load time: {t1 - t0:.4f} sec")

# warmup
print("\n[2] WARMUP...")
warm_df = json_to_df([BASE_TX])
for _ in range(5):
    model.predict(warm_df)

# single inference benchmark
print("\n[3] SINGLE REQUEST BENCH (1000 runs)...")
times = []

for _ in range(1000):
    df = json_to_df([BASE_TX])
    t0 = time.perf_counter()
    model.predict_proba(df)
    times.append(time.perf_counter() - t0)

times = np.array(times)

print(f"avg  : {times.mean()*1000:.2f} ms")
print(f"p95  : {np.percentile(times,95)*1000:.2f} ms")
print(f"max  : {times.max()*1000:.2f} ms")

# batch benchmark
print("\n[4] BATCH 100 BENCH...")
batch = make_batch(100)
df_batch = json_to_df(batch)

t0 = time.perf_counter()
model.predict_proba(df_batch)
dt = time.perf_counter() - t0

print(f"Batch 100 time: {dt:.3f} sec")

# verdict
print("\n[5] VERDICT")
if np.percentile(times,95) < 0.2:
    print("✅ /predict проходит по latency (<200ms p95)")
else:
    print("❌ /predict НЕ проходит по latency")

if dt < 5:
    print("✅ /predict/batch проходит (<5s)")
else:
    print("❌ /predict/batch НЕ проходит")