import os
import joblib
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
from typing import Optional

# ── Paths ─────────────────────────────────────────────────────────────────────
MODELS_DIR = os.path.join(os.path.dirname(__file__), "..", "models")

MODEL_PATHS = {
    "logistic_regression": os.path.join(MODELS_DIR, "logistic_regression.pkl"),
    "random_forest":       os.path.join(MODELS_DIR, "random_forest.pkl"),
    "isolation_forest":    os.path.join(MODELS_DIR, "isolation_forest.pkl"),
}

# Feature order expected by all models (must match training pipeline in 01_data_prep.ipynb)
FEATURE_COLUMNS = [
    "V1",  "V2",  "V3",  "V4",  "V5",  "V6",  "V7",
    "V8",  "V9",  "V10", "V11", "V12", "V13", "V14",
    "V15", "V16", "V17", "V18", "V19", "V20", "V21",
    "V22", "V23", "V24", "V25", "V26", "V27", "V28",
    "Amount_scaled", "Time_scaled",
]

# ── Module-level singletons (loaded once at startup) ──────────────────────────
_models: dict = {}
_scaler_amount: Optional[StandardScaler] = None
_scaler_time:   Optional[StandardScaler] = None


def _build_scalers() -> None:
    """
    Load the fitted StandardScaler objects saved by 01_data_prep.ipynb.
    These ensure Amount and Time are scaled at runtime exactly the same
    way they were scaled during training.
    """
    global _scaler_amount, _scaler_time

    scaler_amount_path = os.path.join(MODELS_DIR, "scaler_amount.pkl")
    scaler_time_path   = os.path.join(MODELS_DIR, "scaler_time.pkl")

    if not os.path.exists(scaler_amount_path) or not os.path.exists(scaler_time_path):
        raise RuntimeError(
            "Scaler files not found in models/. "
            "Please re-run 01_data_prep.ipynb to generate scaler_amount.pkl and scaler_time.pkl."
        )

    _scaler_amount = joblib.load(scaler_amount_path)
    _scaler_time   = joblib.load(scaler_time_path)
    print(f"[model_loader] Loaded scaler_amount from {scaler_amount_path}")
    print(f"[model_loader] Loaded scaler_time   from {scaler_time_path}")


def load_models() -> None:
    """Load all available models from the models/ directory into memory."""
    global _models

    _build_scalers()

    for name, path in MODEL_PATHS.items():
        if os.path.exists(path):
            _models[name] = joblib.load(path)
            print(f"[model_loader] Loaded '{name}' from {path}")
        else:
            print(f"[model_loader] WARNING: model file not found — {path}")

    if not _models:
        raise RuntimeError(
            "No models were loaded. Make sure .pkl files exist in the models/ directory."
        )


def get_model(name: str):
    """Return a loaded model by name. Raises KeyError if not found."""
    if name not in _models:
        available = list(_models.keys())
        raise KeyError(f"Model '{name}' not available. Loaded models: {available}")
    return _models[name]


def get_available_models() -> list[str]:
    """Return the names of all successfully loaded models."""
    return list(_models.keys())


def preprocess_transaction(transaction_dict: dict) -> pd.DataFrame:
    """
    Convert a raw transaction dict (from TransactionInput) into a DataFrame
    that matches the feature space the models were trained on.

    Steps:
      1. Scale 'Amount' and 'Time' with StandardScaler (mirrors 01_data_prep.ipynb)
      2. Drop the original 'Amount' and 'Time' columns
      3. Return a single-row DataFrame with columns in FEATURE_COLUMNS order
    """
    data = dict(transaction_dict)

    # Scale Amount and Time
    amount_scaled = _scaler_amount.transform([[data.pop("Amount")]])[0][0]
    time_scaled   = _scaler_time.transform([[data.pop("Time")]])[0][0]

    data["Amount_scaled"] = amount_scaled
    data["Time_scaled"]   = time_scaled

    # Build DataFrame with correct column order
    df = pd.DataFrame([data])[FEATURE_COLUMNS]
    return df


def predict(model_name: str, transaction_dict: dict) -> dict:
    """
    Run end-to-end prediction for a single transaction.

    For supervised models (logistic_regression, random_forest):
      - Uses predict_proba to return a fraud probability

    For unsupervised models (isolation_forest):
      - Uses decision_function to return an anomaly score
      - Negative scores indicate anomalies (fraud)
    """
    model = get_model(model_name)
    X = preprocess_transaction(transaction_dict)

    if model_name == "isolation_forest":
        raw_prediction = model.predict(X)[0]
        anomaly_score  = float(model.score_samples(X)[0])
        is_fraud       = raw_prediction == -1
        fraud_proba    = float(1 / (1 + np.exp(5 * anomaly_score)))
        confidence     = fraud_proba if is_fraud else 1.0 - fraud_proba
    else:
        anomaly_score = None
        fraud_proba   = float(model.predict_proba(X)[0][1])
        is_fraud      = fraud_proba >= 0.5
        confidence    = fraud_proba if is_fraud else 1.0 - fraud_proba

    if fraud_proba < 0.3:
        risk_level = "LOW"
    elif fraud_proba < 0.6:
        risk_level = "MEDIUM"
    else:
        risk_level = "HIGH"

    result = {
        "is_fraud":          is_fraud,
        "fraud_probability": round(fraud_proba, 6),
        "risk_level":        risk_level,
        "confidence":        round(confidence, 6),
    }

    if anomaly_score is not None:
        result["anomaly_score"] = round(anomaly_score, 6)

    return result