"""
Taxy ML — Inference
===================
Kode sederhana untuk proses inferensi model: memuat model .keras + artefak
preprocessing, lalu memprediksi risiko ketidakpatuhan pajak dari profil pengguna.

Dipakai sebagai modul oleh serve.py, dan bisa dijalankan langsung untuk demo:
    python inference.py
"""

import os
import numpy as np

# keras/tensorflow di-import lazily agar startup ringan saat hanya dipakai modul
from preprocessing import (
    build_raw_features,
    load_preprocessing,
    CLASS_NAMES,
)

HERE = os.path.dirname(os.path.abspath(__file__))
ARTIFACTS = os.path.join(HERE, "artifacts")
MODEL_PATH = os.path.join(ARTIFACTS, "taxy_risk_model.keras")
PREP_PATH = os.path.join(ARTIFACTS, "preprocessing.json")

_model = None
_prep = None


def load():
    """Muat model + preprocessing (sekali)."""
    global _model, _prep
    if _model is None:
        import keras
        # custom objects teregistrasi via @register_keras_serializable di import
        import custom_components  # noqa: F401  (registrasi FeatureGatingLayer dll)
        _model = keras.models.load_model(MODEL_PATH)
        _prep = load_preprocessing(PREP_PATH)
    return _model, _prep


def predict(record):
    """Prediksi risiko dari satu record profil pengguna (dict).

    Returns dict: risk_index, risk_category, confidence, probabilities.
    """
    model, prep = load()
    x = build_raw_features(record).reshape(1, -1)
    x = prep["scaler"].transform(x)
    reg, cls = model.predict(x, verbose=0)
    risk_index = float(reg[0][0])
    probs = cls[0]
    cls_idx = int(np.argmax(probs))
    return {
        "risk_index": round(risk_index, 4),
        "risk_category": CLASS_NAMES[cls_idx],
        "risk_category_id": cls_idx,
        "confidence": round(float(probs[cls_idx]), 4),
        "probabilities": {CLASS_NAMES[i]: round(float(p), 4) for i, p in enumerate(probs)},
    }


def predict_batch(records):
    model, prep = load()
    X = np.stack([build_raw_features(r) for r in records])
    X = prep["scaler"].transform(X)
    reg, cls = model.predict(X, verbose=0)
    out = []
    for i in range(len(records)):
        probs = cls[i]
        cls_idx = int(np.argmax(probs))
        out.append({
            "risk_index": round(float(reg[i][0]), 4),
            "risk_category": CLASS_NAMES[cls_idx],
            "confidence": round(float(probs[cls_idx]), 4),
        })
    return out


if __name__ == "__main__":
    samples = [
        {"usia": 28, "gaji_bulanan": 8_000_000, "skor_literasi_pajak": 80,
         "punya_npwp": 1, "aktif_taxy": 1, "status_ptkp": "TK/0",
         "pekerjaan": "Senior Staff"},
        {"usia": 45, "gaji_bulanan": 26_000_000, "skor_literasi_pajak": 25,
         "punya_npwp": 0, "aktif_taxy": 0, "status_ptkp": "K/3",
         "pekerjaan": "Manager"},
        {"usia": 22, "gaji_bulanan": 4_500_000, "skor_literasi_pajak": 55,
         "punya_npwp": 1, "aktif_taxy": 0, "status_ptkp": "TK/0",
         "pekerjaan": "Fresh Graduate"},
    ]
    print("=== DEMO INFERENCE TAXY RISK PROFILER ===\n")
    for s in samples:
        r = predict(s)
        print(f"Profil : usia {s['usia']}, gaji {s['gaji_bulanan']:,}, "
              f"NPWP={s['punya_npwp']}, literasi={s['skor_literasi_pajak']}")
        print(f"  -> Risiko: {r['risk_category']} "
              f"(index={r['risk_index']}, confidence={r['confidence']})")
        print(f"     prob: {r['probabilities']}\n")