"""Purchase Propensity + RFM: inference code of the Hugging Face model repo shalev396/purchase-propensity.

Single source of truth for the input schema, the preprocessing, the PyTorch MLP architecture and
inference. The training code (../training), the Space and the Inference Endpoint handler all import
this file, so evaluation and serving run exactly the same code path.

A website session is described by 23 binary flags (did the visitor add to basket, see the checkout,
sign in, use a mobile device, ...). Two models score P(the session ends in an order):

    logistic_regression   scikit-learn LogisticRegression(class_weight="balanced")   24 parameters
    pytorch_mlp           PyTorch MLP 23 -> 64 -> 32 -> 1 (ReLU + dropout)            3,649 parameters

Files next to this module (written by training/src/export.py):
    config.json               feature names, MLP architecture, both models + tuned thresholds, default model
    model.safetensors         PyTorch MLP weights (PyTorchModelHubMixin)
    logreg.joblib             fitted LogisticRegression
    scaler.joblib             StandardScaler fit on the training split (used by both models)
    rfm_segment_stats.json    RFM customer segments of the UCI Online Retail data + quintile edges

    import model
    predictor = model.load("path/to/this/repo", device="cpu")
    predictor.predict(["basket_add_detail", "saw_checkout", "checked_delivery_detail"])
    # -> {"buy": 0.99, "no_buy": 0.01}
"""
from __future__ import annotations

import json
import os
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
import torch
from huggingface_hub import PyTorchModelHubMixin
from torch import nn

REPO_ID = "shalev396/purchase-propensity"
FRAMEWORK = "sklearn+pytorch"
TARGET = "ordered"  # label column of the Kaggle data (optional in files you score)

# The 23 binary session flags of the Kaggle "customer propensity to purchase" data, in model input order.
FEATURES = [
    "basket_icon_click", "basket_add_list", "basket_add_detail", "sort_by", "image_picker",
    "account_page_click", "promo_banner_click", "detail_wishlist_add", "list_size_dropdown",
    "closed_minibasket_click", "checked_delivery_detail", "checked_returns_detail", "sign_in",
    "saw_checkout", "saw_sizecharts", "saw_delivery", "saw_account_upgrade", "saw_homepage",
    "device_mobile", "device_computer", "device_tablet", "returning_user", "loc_uk",
]
DEVICE_FEATURES = ["device_mobile", "device_computer", "device_tablet"]
VISITOR_FEATURES = ["returning_user", "loc_uk"]

MODEL_NAMES = ("logistic_regression", "pytorch_mlp")
WEIGHTS_FILE = "model.safetensors"
RFM_FILE = "rfm_segment_stats.json"


def cuda_available() -> bool:
    """The MLP runs on CUDA when present; the logistic regression always runs on CPU."""
    return torch.cuda.is_available()


# --------------------------------------------------------------------------- architecture
class PropensityMLP(
    nn.Module,
    PyTorchModelHubMixin,
    tags=["ml-lab", "tabular-classification", "propensity"],
    repo_url="https://github.com/shalev396/ml-lab",
    pipeline_tag="tabular-classification",
    license="mit",
):
    """`Linear -> ReLU -> Dropout` blocks for every entry of `hidden_units`, then one logit.

    Train with `BCEWithLogitsLoss`; `predict_proba` applies the sigmoid. Inputs are the
    standardized flags (`scaler.joblib`), in `feature_names` order.
    """

    def __init__(self, feature_names: list[str] | None = None, hidden_units: list[int] | None = None,
                 dropout: float = 0.2):
        super().__init__()
        self.feature_names = list(feature_names or FEATURES)
        self.hidden_units = [int(u) for u in (hidden_units or [64, 32])]
        layers: list[nn.Module] = []
        prev = len(self.feature_names)
        for units in self.hidden_units:
            layers += [nn.Linear(prev, units), nn.ReLU(), nn.Dropout(dropout)]
            prev = units
        layers.append(nn.Linear(prev, 1))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """(N, n_features) standardized flags -> (N, 1) logits."""
        return self.net(x)


# --------------------------------------------------------------------------- input handling
def _flag(name: str, value) -> float:
    if isinstance(value, str):
        value = value.strip().lower()
        value = {"true": 1, "false": 0, "yes": 1, "no": 0}.get(value, value)
    try:
        number = float(value)
    except (TypeError, ValueError):
        raise ValueError(f"feature {name!r} must be 0 or 1, got {value!r}") from None
    if number not in (0.0, 1.0):
        raise ValueError(f"feature {name!r} must be 0 or 1, got {value!r}")
    return number


def encode(features, feature_names: list[str] = FEATURES) -> np.ndarray:
    """One session -> (1, n_features) float array of 0/1 flags.

    `features` is either a dict {feature_name: 0/1} (missing names count as 0) or a list of the
    names of the flags that are ON. Unknown names raise ValueError.
    """
    if features is None:
        features = []
    if isinstance(features, dict):
        items = {str(k): _flag(str(k), v) for k, v in features.items()}
    elif isinstance(features, (list, tuple, set)):
        items = {str(name): 1.0 for name in features}
    else:
        raise ValueError("features must be a dict {name: 0/1} or a list of active feature names")
    unknown = sorted(set(items) - set(feature_names))
    if unknown:
        raise ValueError(f"unknown feature(s) {unknown}; valid names: {feature_names}")
    return np.array([[items.get(name, 0.0) for name in feature_names]], dtype=np.float64)


def read_sessions(data, feature_names: list[str] = FEATURES) -> pd.DataFrame:
    """DataFrame or CSV path -> validated 0/1 frame with the feature columns (extra columns are ignored)."""
    frame = pd.read_csv(data) if isinstance(data, (str, os.PathLike)) else pd.DataFrame(data)
    missing = [c for c in feature_names if c not in frame.columns]
    if missing:
        raise ValueError(f"missing required column(s): {', '.join(missing)}")
    values = frame[feature_names].apply(pd.to_numeric, errors="coerce").reset_index(drop=True)
    if values.empty:
        raise ValueError("no sessions to score")
    bad = np.flatnonzero((~values.isin([0, 1])).any(axis=1).to_numpy())
    if bad.size:
        raise ValueError(f"values must be 0 or 1; bad row(s) {bad[:10].tolist()} (0-based)")
    return values.astype(np.float64)


# --------------------------------------------------------------------------- RFM helpers
def rfm_scores(recency: float, frequency: float, monetary: float, edges: dict) -> tuple[int, int, int]:
    """Quintile scores 1..5 for one customer, from the training quintile edges.

    Recency: fewer days since the last order = higher score. Frequency / monetary: more = higher.
    `edges[k]` holds the 4 inner cut points of the quintiles (upper bound of scores 1..4).
    """
    def score(value, cuts, reverse=False):
        s = int(np.searchsorted(np.asarray(cuts, dtype=float), float(value), side="left")) + 1
        return 6 - s if reverse else s
    return (score(recency, edges["recency"], reverse=True), score(frequency, edges["frequency"]),
            score(monetary, edges["monetary"]))


def rfm_segment(r_score: int, f_score: int, segment_map: list[list[str]]) -> str:
    """Standard R x F segment grid (`segment_map` = [[regex on "<R><F>", segment name], ...])."""
    import re

    key = f"{r_score}{f_score}"
    for pattern, name in segment_map:
        if re.fullmatch(pattern, key):
            return name
    return "Other"


# --------------------------------------------------------------------------- inference
class Predictor:
    """Scores website sessions with the logistic regression or the MLP (default: the best on validation)."""

    def __init__(self, model_dir: str | Path, device: str = "cpu"):
        self.model_dir = Path(model_dir)
        self.device = torch.device(device)
        self.config = json.loads(self._require("config.json").read_text(encoding="utf-8"))
        self.feature_names = list(self.config["feature_names"])
        self.best = self.config["best"]
        self.scaler = joblib.load(self._require(self.config["scaler"]))
        self.logreg = joblib.load(self._require(self.config["models"]["logistic_regression"]["file"]))
        self._require(WEIGHTS_FILE)
        self.mlp = PropensityMLP.from_pretrained(self.model_dir, map_location="cpu", strict=True)
        self.mlp.to(self.device).eval()
        rfm_path = self.model_dir / RFM_FILE
        self.rfm = json.loads(rfm_path.read_text(encoding="utf-8")) if rfm_path.is_file() else None

    def _require(self, filename: str) -> Path:
        path = self.model_dir / filename
        if not path.is_file():
            raise FileNotFoundError(f"{path} not found: train and export the model first "
                                    f"(training/notebook.ipynb) or snapshot_download('{REPO_ID}')")
        return path

    # ---- model selection
    @property
    def choices(self) -> list[str]:
        """Valid values for `model`: "best" plus every exported model name."""
        return ["best", *self.config["models"]]

    def resolve(self, model: str | None = "best") -> str:
        name = self.best if model in (None, "", "best") else model
        if name not in self.config["models"]:
            raise ValueError(f"unknown model {model!r}; choose one of {self.choices}")
        return name

    def threshold(self, model: str | None = "best") -> float:
        """Decision threshold tuned on the validation split (max F1) for this model."""
        return float(self.config["models"][self.resolve(model)]["threshold"])

    def runs_on(self, model: str | None = "best") -> str:
        """"gpu" when this model runs on CUDA, else "cpu" (the logistic regression is CPU-only)."""
        return "gpu" if self.resolve(model) == "pytorch_mlp" and self.device.type == "cuda" else "cpu"

    # ---- scoring
    def probabilities(self, x: np.ndarray, model: str | None = "best") -> np.ndarray:
        """(N, n_features) raw 0/1 flags -> P(order) per row."""
        scaled = self.scaler.transform(np.asarray(x, dtype=np.float64))
        if self.resolve(model) == "logistic_regression":
            return self.logreg.predict_proba(scaled)[:, 1]
        with torch.inference_mode():
            t = torch.as_tensor(scaled, dtype=torch.float32, device=self.device)
            return torch.sigmoid(self.mlp(t)).float().cpu().numpy().ravel().astype(np.float64)

    def predict(self, features, model: str = "best"):
        """One session -> {"buy": p, "no_buy": 1 - p}.

        `features`: a dict {feature_name: 0/1} or a list of active feature names. A list of dicts
        scores several sessions and returns a list of results.
        """
        if isinstance(features, list) and features and all(isinstance(f, dict) for f in features):
            return [self.predict(f, model) for f in features]
        p = float(self.probabilities(encode(features, self.feature_names), model)[0])
        return {"buy": p, "no_buy": 1.0 - p}

    def explain(self, features, model: str = "best") -> dict[str, float]:
        """Occlusion: how much P(buy) drops when each active flag is switched off (sorted by size)."""
        x = encode(features, self.feature_names)
        active = [i for i, v in enumerate(x[0]) if v == 1.0]
        if not active:
            return {}
        variants = np.repeat(x, len(active) + 1, axis=0)
        for row, col in enumerate(active, start=1):
            variants[row, col] = 0.0
        prob = self.probabilities(variants, model)
        deltas = {self.feature_names[col]: float(prob[0] - prob[row]) for row, col in enumerate(active, start=1)}
        return dict(sorted(deltas.items(), key=lambda kv: abs(kv[1]), reverse=True))

    def predict_batch(self, data, model: str = "best") -> dict:
        """Score every row of a DataFrame / CSV with the 23 flag columns.

        Returns {"model", "threshold", "n_rows", "n_buyers",
                 "rows": [{"row": 0-based index, "buy_probability": p, "will_buy": p >= threshold}]}.
        """
        name = self.resolve(model)
        thr = self.threshold(name)
        prob = self.probabilities(read_sessions(data, self.feature_names).to_numpy(), name)
        flags = prob >= thr
        return {
            "model": name,
            "threshold": thr,
            "n_rows": int(prob.size),
            "n_buyers": int(flags.sum()),
            "rows": [{"row": i, "buy_probability": round(float(p), 6), "will_buy": bool(f)}
                     for i, (p, f) in enumerate(zip(prob, flags))],
        }

    # ---- RFM segmentation (UCI Online Retail)
    def segments(self) -> list[dict]:
        """RFM segment table: customers, share, average recency/frequency/monetary, revenue."""
        if self.rfm is None:
            raise FileNotFoundError(f"{self.model_dir / RFM_FILE} not found")
        return list(self.rfm["segments"])

    def rfm_segment(self, recency_days: float, frequency: float, monetary: float) -> dict:
        """Place one customer (days since last order, number of orders, total spend in GBP) on the grid."""
        if self.rfm is None:
            raise FileNotFoundError(f"{self.model_dir / RFM_FILE} not found")
        if min(recency_days, frequency, monetary) < 0:
            raise ValueError("recency, frequency and monetary must be >= 0")
        r, f, m = rfm_scores(recency_days, frequency, monetary, self.rfm["edges"])
        name = rfm_segment(r, f, self.rfm["segment_map"])
        stats = next((s for s in self.rfm["segments"] if s["segment"] == name), None)
        return {"r_score": r, "f_score": f, "m_score": m, "rfm_score": f"{r}{f}{m}", "segment": name,
                "segment_stats": stats}


def load(model_dir: str | Path, device: str = "cpu") -> Predictor:
    return Predictor(model_dir, device)