inference.py

"""
Stacking classifier inference for SBAN multi-representation malware dataset classification.

Loads sban_weighted_stacking_model.joblib and predicts which source dataset
(bodmas, dike, malwarebazaar, sorel20m) each sample belongs to.
"""

from __future__ import annotations

import argparse
import math
import re
from collections import Counter
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
from scipy.sparse import csr_matrix, hstack
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    f1_score,
)

PROJECT_DIR = Path(__file__).resolve().parent
DEFAULT_MODEL_PATH = PROJECT_DIR / "sban_weighted_stacking_model.joblib"

ID_COLUMN = "ID"
LABEL_COLUMN = "dataset_name"

REPRESENTATION_COLUMNS = {
    "asm": "assembly_code",
    "binary": "binary_code",
    "source": "source_code",
    "nld": "NLD",
}

NUMERIC_FEATURE_NAMES = [
    "char_count",
    "line_count",
    "token_count",
    "avg_line_length",
    "unique_token_ratio",
    "char_entropy",
]


def binary_to_byte_tokens(text: str) -> str:
    tokens: list[str] = []

    for line in text.splitlines():
        line = re.sub(r"\s+", "", line)

        if not line or re.fullmatch(r"[0-9a-fA-F]+", line) is None:
            continue

        byte_count = len(line) // 2
        tokens.extend(line[i * 2 : (i + 1) * 2].lower() for i in range(byte_count))
        tokens.append("instsep")

    return " ".join(tokens)


def calculate_char_entropy(text: str) -> float:
    if not text:
        return 0.0

    counts = Counter(text)
    length = len(text)

    return -sum(
        (count / length) * math.log2(count / length) for count in counts.values()
    )


def extract_numeric_features(text: str) -> list[float]:
    lines = text.splitlines()
    tokens = text.split()

    char_count = len(text)
    line_count = len(lines)
    token_count = len(tokens)

    avg_line_length = (
        sum(len(line) for line in lines) / line_count if line_count else 0.0
    )

    unique_token_ratio = (
        len(set(tokens)) / token_count if token_count else 0.0
    )

    return [
        float(char_count),
        float(line_count),
        float(token_count),
        float(avg_line_length),
        float(unique_token_ratio),
        calculate_char_entropy(text),
    ]


def build_numeric_features(series: pd.Series) -> np.ndarray:
    return np.asarray(
        [extract_numeric_features(text) for text in series.astype(str)],
        dtype=np.float32,
    )


def build_inference_features(
    representation_name: str,
    text_series: pd.Series,
    representation_bundle: dict,
) -> tuple[csr_matrix, np.ndarray]:
    texts = text_series.fillna("").astype(str).tolist()

    if representation_name == "binary":
        tfidf_texts = [binary_to_byte_tokens(text) for text in texts]
    else:
        tfidf_texts = texts

    vectorizer = representation_bundle["vectorizer"]
    numeric_scaler = representation_bundle["numeric_scaler"]

    x_tfidf = vectorizer.transform(tfidf_texts)
    x_numeric = build_numeric_features(text_series)
    x_numeric_scaled = numeric_scaler.transform(x_numeric)

    x_full = hstack(
        [x_tfidf, csr_matrix(x_numeric_scaled)],
        format="csr",
        dtype=np.float32,
    )

    return x_full, np.asarray(x_numeric_scaled, dtype=np.float32)


class StackingPredictor:
    """Wrapper around the exported joblib inference bundle."""

    def __init__(self, model_path: Path | str = DEFAULT_MODEL_PATH) -> None:
        self.model_path = Path(model_path)
        self.model_bundle = joblib.load(self.model_path)
        self.representation_order = self.model_bundle["representation_order"]
        self.representation_columns = self.model_bundle["representation_columns"]
        self.label_encoder = self.model_bundle["label_encoder"]
        self.meta_model = self.model_bundle["meta_model"]
        self.representation_bundles = self.model_bundle["representations"]

    @property
    def classes(self) -> list[str]:
        return self.label_encoder.classes_.tolist()

    def predict(self, df: pd.DataFrame) -> pd.DataFrame:
        missing = [
            col
            for rep in self.representation_order
            if (col := self.representation_columns[rep]) not in df.columns
        ]
        if missing:
            raise ValueError(f"Missing required columns: {missing}")

        meta_features: list[np.ndarray] = []

        for representation_name in self.representation_order:
            bundle = self.representation_bundles[representation_name]

            x_full, x_numeric = build_inference_features(
                representation_name,
                df[self.representation_columns[representation_name]],
                bundle,
            )

            x_selected = x_full[:, bundle["selected_indices"]]
            decision_scores = bundle["base_model"].decision_function(x_selected)

            meta_features.append(decision_scores)
            meta_features.append(x_numeric)

        x_meta = np.hstack(meta_features)

        encoded_predictions = self.meta_model.predict(x_meta)
        prediction_probabilities = self.meta_model.predict_proba(x_meta)

        result = df.copy()
        result["prediction"] = self.label_encoder.inverse_transform(
            encoded_predictions
        )

        for index, class_name in enumerate(self.label_encoder.classes_):
            result[f"prob_{class_name}"] = prediction_probabilities[:, index]

        return result


def load_predictor(model_path: Path | str = DEFAULT_MODEL_PATH) -> StackingPredictor:
    return StackingPredictor(model_path)


def evaluate_predictions(
    predictions: pd.DataFrame,
    label_column: str = LABEL_COLUMN,
) -> None:
    if label_column not in predictions.columns:
        print(f"Skip evaluation: column {label_column!r} not in input.")
        return

    y_true = predictions[label_column]
    y_pred = predictions["prediction"]

    print(f"Accuracy    : {accuracy_score(y_true, y_pred):.6f}")
    print(f"Macro F1    : {f1_score(y_true, y_pred, average='macro'):.6f}")
    print(f"Weighted F1 : {f1_score(y_true, y_pred, average='weighted'):.6f}")
    print("\nClassification Report\n")
    print(
        classification_report(
            y_true,
            y_pred,
            digits=4,
            target_names=sorted(y_true.unique()),
        )
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run SBAN weighted stacking inference on a parquet file.",
    )
    parser.add_argument(
        "--model-path",
        type=Path,
        default=DEFAULT_MODEL_PATH,
        help="Path to sban_weighted_stacking_model.joblib",
    )
    parser.add_argument(
        "--input",
        type=Path,
        required=True,
        help="Input parquet with representation columns",
    )
    parser.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Output parquet with prediction and prob_* columns",
    )
    parser.add_argument(
        "--evaluate",
        action="store_true",
        help="Print metrics when dataset_name column is present",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if not args.model_path.is_file():
        raise FileNotFoundError(f"Model not found: {args.model_path}")
    if not args.input.is_file():
        raise FileNotFoundError(f"Input not found: {args.input}")

    predictor = load_predictor(args.model_path)
    df = pd.read_parquet(args.input)
    predictions = predictor.predict(df)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    predictions.to_parquet(args.output, index=False)

    print(f"Wrote {len(predictions)} rows to {args.output}")

    if args.evaluate:
        evaluate_predictions(predictions)


if __name__ == "__main__":
    main()