src/services/predictor_svc.py

"""
TradeFlow AI — XGBoost Rejection Predictor (Phase 3, Step 3.2)

PRD §13 — Trains on submission_outcomes table.
Features: CRS pillars, HS confidence, company historical rate, duty value.
Target: binary (0=accepted, 1=rejected).

The model artifact is stored in MinIO at:
  models/rejection_predictor/model_v{version}.joblib
"""

from __future__ import annotations

import io
from pathlib import Path
from typing import Any

import structlog

try:
    import numpy as np
except Exception:  # pragma: no cover - optional in lightweight test environments
    np = None

try:
    import joblib
except Exception:  # pragma: no cover - optional in lightweight test environments
    joblib = None

try:
    from sklearn.pipeline import Pipeline
except Exception:  # pragma: no cover - optional in lightweight test environments
    Pipeline = object

from ..config import settings

log = structlog.get_logger()

FEATURE_NAMES = [
    "doc_quality_score",
    "completeness_score",
    "consistency_score",
    "historical_rate",
    "hs_confidence",
    "cif_value_usd",
    "package_count",
    "gross_weight_kg",
]

MODEL_LOCAL_PATH = Path("/tmp/rejection_predictor.joblib")


class RejectionPredictor:
    """
    XGBoost binary classifier for customs declaration rejection prediction.
    """

    def __init__(self) -> None:
        self._pipeline: Pipeline | None = None
        self._model_version: str = "0"
        self._auc: float | None = None
        self._load_attempted = False

    def _load_from_minio(self) -> bool:
        """Download latest model from MinIO."""
        if joblib is None:
            log.warning("joblib is not installed; using heuristic predictor fallback")
            return False

        try:
            import boto3
            from botocore.client import Config

            s3 = boto3.client(
                "s3",
                endpoint_url=f"http://{settings.MINIO_ENDPOINT}",
                aws_access_key_id=settings.MINIO_ACCESS_KEY,
                aws_secret_access_key=settings.MINIO_SECRET_KEY,
                config=Config(signature_version="s3v4"),
                region_name="us-east-1",
            )
            # List versions, pick latest
            response = s3.list_objects_v2(
                Bucket=settings.STORAGE_BUCKET_NAME,
                Prefix="models/rejection_predictor/",
            )
            objects = sorted(
                response.get("Contents", []),
                key=lambda x: x["LastModified"],
                reverse=True,
            )
            if not objects:
                log.warning("No model found in MinIO — using heuristic fallback")
                return False

            latest_key = objects[0]["Key"]
            self._model_version = latest_key.split("model_v")[1].replace(".joblib", "")

            buf = io.BytesIO()
            s3.download_fileobj(settings.STORAGE_BUCKET_NAME, latest_key, buf)
            buf.seek(0)
            self._pipeline = joblib.load(buf)
            log.info("Model loaded from MinIO", version=self._model_version, key=latest_key)
            return True
        except Exception as exc:
            log.error("Failed to load model from MinIO", error=str(exc))
            return False

    def load(self) -> None:
        """Load model — try MinIO, fallback to local cache."""
        self._load_attempted = True
        if not self._load_from_minio():
            if MODEL_LOCAL_PATH.exists():
                if joblib is None:
                    log.warning("Local model exists but joblib is unavailable")
                    return
                self._pipeline = joblib.load(MODEL_LOCAL_PATH)
                log.info("Model loaded from local cache")
            else:
                log.warning("No model available — predictions will use heuristic")

    def predict_proba(self, features: dict) -> float:
        """
        Returns probability of rejection (0.0–1.0).
        If model is not loaded, falls back to heuristic.
        """
        if np is not None:
            feature_vec = np.array(
                [[features.get(f, 0.0) for f in FEATURE_NAMES]], dtype=np.float32
            )
        else:
            feature_vec = [[features.get(f, 0.0) for f in FEATURE_NAMES]]

        if self._pipeline is None and not self._load_attempted:
            self.load()

        if self._pipeline is not None:
            try:
                prob = float(self._pipeline.predict_proba(feature_vec)[0][1])
                log.debug("XGBoost prediction", prob=prob, version=self._model_version)
                return prob
            except Exception as exc:
                log.error("XGBoost inference error", error=str(exc))

        # Heuristic fallback: 1 - average of available scores
        known = [v for k, v in features.items() if "score" in k or "confidence" in k]
        return round(1.0 - (sum(known) / len(known)), 4) if known else 0.5

    def train_and_upload(
        self,
        X: Any,
        y: Any,
        version: str,
        *,
        current_auc: float | None = None,
    ) -> dict:
        """
        Train a new XGBoost model and upload to MinIO.
        Called by the Celery `retrain_predictor` task when ≥100 new samples.
        """
        if np is None:
            raise RuntimeError("numpy is required to train the XGBoost model")
        if joblib is None:
            raise RuntimeError("joblib is required to train and persist the XGBoost model")

        import xgboost as xgb
        from sklearn.metrics import roc_auc_score
        from sklearn.model_selection import train_test_split
        from sklearn.pipeline import Pipeline
        from sklearn.preprocessing import StandardScaler

        if len(set(y.tolist())) < 2:
            raise ValueError("Cannot train rejection predictor with only one outcome class")

        pipeline = Pipeline([
            ("scaler", StandardScaler()),
            ("xgb", xgb.XGBClassifier(
                n_estimators=300,
                max_depth=6,
                learning_rate=0.05,
                subsample=0.8,
                colsample_bytree=0.8,
                use_label_encoder=False,
                eval_metric="logloss",
                random_state=42,
            )),
        ])
        stratify = y if min(np.bincount(y.astype(int))) >= 2 else None
        X_train, X_valid, y_train, y_valid = train_test_split(
            X,
            y,
            test_size=0.2,
            random_state=42,
            stratify=stratify,
        )
        pipeline.fit(X_train, y_train)
        valid_prob = pipeline.predict_proba(X_valid)[:, 1]
        auc = float(roc_auc_score(y_valid, valid_prob)) if len(set(y_valid.tolist())) > 1 else 0.5

        baseline_auc = current_auc if current_auc is not None else self._auc
        if baseline_auc is not None and auc < (baseline_auc - settings.RETRAIN_MAX_AUC_DROP):
            log.warning(
                "Rejected candidate model because AUC regressed",
                candidate_auc=auc,
                baseline_auc=baseline_auc,
            )
            return {"promoted": False, "auc": auc, "baseline_auc": baseline_auc}

        buf = io.BytesIO()
        joblib.dump(pipeline, buf)
        buf.seek(0)

        # Upload to MinIO
        try:
            import boto3
            from botocore.client import Config

            s3 = boto3.client(
                "s3",
                endpoint_url=f"http://{settings.MINIO_ENDPOINT}",
                aws_access_key_id=settings.MINIO_ACCESS_KEY,
                aws_secret_access_key=settings.MINIO_SECRET_KEY,
                config=Config(signature_version="s3v4"),
                region_name="us-east-1",
            )
            key = f"models/rejection_predictor/model_v{version}.joblib"
            s3.upload_fileobj(buf, settings.STORAGE_BUCKET_NAME, key)
            log.info("New model uploaded to MinIO", version=version, key=key)
            self._pipeline = pipeline
            self._model_version = version
            self._auc = auc
            return {"promoted": True, "auc": auc, "version": version, "key": key}
        except Exception as exc:
            log.error("Failed to upload model to MinIO", error=str(exc))
            raise


# ── Singleton — loaded once per Celery worker startup ─────────────────────────
rejection_predictor = RejectionPredictor()