utils/LocalModelStability.py

import numpy as np
import torch


class LocalModelStability:
    def __init__(self, model, reference_data, feature_specs, device="cpu"):
        """
        model           : PyTorch model; model(x, apply_activation=False) returns logits
        reference_data  : numpy array (n_rows, n_columns), one-hot encoded
        feature_specs   : list of dicts, each with:
                            - 'name': str
                            - 'type': 'numerical' | 'ordinal_group' | 'categorical_group'
                            - 'columns': list of column indices
                            - 'decoded_values': list of ints (ordinal_group only)
        device          : 'cpu' or 'cuda'
        """
        self.model = model.to(device).eval()
        self.device = device
        self.reference_data = np.asarray(reference_data, dtype=np.float32)
        self.feature_specs = feature_specs
        self._spec_by_name = {spec["name"]: spec for spec in feature_specs}

        # Pre-decode ordinal groups
        self._decoded_ordinals = {}
        for spec in feature_specs:
            if spec["type"] == "ordinal_group":
                cols = spec["columns"]
                values = np.asarray(spec["decoded_values"])
                block = self.reference_data[:, cols]                # (n_rows, group_size)
                active_col = np.argmax(block, axis=1)               # (n_rows,)
                has_active = block.sum(axis=1) > 0                  # (n_rows,) bool
                decoded = np.where(has_active, values[active_col], np.iinfo(np.int64).min)
                self._decoded_ordinals[spec["name"]] = decoded

        # Pre-compute reference predictions (batched)
        self._reference_predictions = self._batch_predict(self.reference_data)

    # ------------------------------------------------------------
    # Prediction
    # ------------------------------------------------------------
    def _batch_predict(self, X, batch_size=1024):
        preds = []
        with torch.no_grad():
            for i in range(0, len(X), batch_size):
                batch = torch.tensor(X[i:i + batch_size], dtype=torch.float32).to(self.device)
                out = self.model(batch, apply_activation=False).cpu().numpy().reshape(-1)
                preds.append(out)
        return np.concatenate(preds)

    def predict(self, x):
        """Return logit(s) for input x (1D or 2D array)."""
        x = np.atleast_2d(np.asarray(x, dtype=np.float32))
        return self._batch_predict(x)

    # ------------------------------------------------------------
    # Neighborhood building
    # ------------------------------------------------------------
    def build_neighborhood(self, x0, feature_name, n_max, rng=None):
        if rng is None:
            rng = np.random.default_rng()

        x0 = np.asarray(x0, dtype=np.float32).reshape(-1)
        spec = self._spec_by_name[feature_name]
        ftype = spec["type"]

        if ftype == "numerical":
            return self._build_numerical(x0, spec, n_max, rng)
        elif ftype == "ordinal_group":
            return self._build_ordinal(x0, spec, n_max, rng)
        elif ftype == "categorical_group":
            return self._build_categorical(x0, spec, n_max, rng)
        else:
            raise ValueError(f"Unknown feature type: {ftype}")

    # ---- Numerical ----
    def _build_numerical(self, x0, spec, n_max, rng):
        col = spec["columns"][0]
        target = x0[col]
        feature_values = self.reference_data[:, col]

        exact_mask = feature_values == target
        exact_idx = np.where(exact_mask)[0]
        n_exact = len(exact_idx)

        if n_exact >= n_max:
            # Take ALL exact matches (no subsampling)
            selected = exact_idx
        else:
            # Fill up to n_max with nearest neighbors
            remaining = n_max - n_exact
            non_exact_idx = np.where(~exact_mask)[0]
            if len(non_exact_idx) == 0:
                selected = exact_idx
            else:
                distances = np.abs(feature_values[non_exact_idx] - target)
                k = min(remaining, len(non_exact_idx))
                nearest = non_exact_idx[np.argpartition(distances, k - 1)[:k]]
                selected = np.concatenate([exact_idx, nearest])

        return self._package_result(selected, n_exact, "ok")
    # ---- Ordinal group ----
    def _build_ordinal(self, x0, spec, n_max, rng):
        cols = spec["columns"]
        values = np.asarray(spec["decoded_values"])
        x0_group = x0[cols]

        if x0_group.sum() == 0:
            return self._empty_result("no_active_category")

        target_value = values[int(np.argmax(x0_group))]
        decoded_ref = self._decoded_ordinals[spec["name"]]

        # Exclude rows with no active category
        valid_mask = decoded_ref != np.iinfo(np.int64).min
        valid_idx = np.where(valid_mask)[0]
        valid_values = decoded_ref[valid_idx]

        exact_mask = valid_values == target_value
        exact_idx = valid_idx[exact_mask]
        n_exact = len(exact_idx)

        if n_exact >= n_max:
            # Take ALL exact matches (no subsampling)
            selected = exact_idx
        else:
            # Exact matches insufficient — fill up to n_max with nearest neighbors
            remaining = n_max - n_exact
            non_exact_idx = valid_idx[~exact_mask]
            if len(non_exact_idx) == 0:
                selected = exact_idx
            else:
                distances = np.abs(decoded_ref[non_exact_idx] - target_value)
                k = min(remaining, len(non_exact_idx))
                nearest = non_exact_idx[np.argpartition(distances, k - 1)[:k]]
                selected = np.concatenate([exact_idx, nearest])

        return self._package_result(selected, n_exact, "ok")

    # ---- Categorical group ----
    def _build_categorical(self, x0, spec, n_max, rng):
        cols = spec["columns"]
        x0_group = x0[cols]

        if x0_group.sum() == 0:
            return self._empty_result("no_active_category")

        active_col_in_group = int(np.argmax(x0_group))
        active_col_global = cols[active_col_in_group]

        match_mask = self.reference_data[:, active_col_global] == 1.0
        match_idx = np.where(match_mask)[0]
        n_matches = len(match_idx)

        if n_matches == 0:
            return self._empty_result("empty")

        selected = match_idx   # take all exact matches

        return self._package_result(selected, n_matches if n_matches < n_max else n_max, "ok")

    # ------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------
    def _package_result(self, indices, n_exact, status):
        return {
            "indices":         indices,
            "predictions":     self._reference_predictions[indices],
            "n_selected":      len(indices),
            "n_exact_matches": int(n_exact),
            "status":          status,
        }

    def _empty_result(self, status):
        return {
            "indices":         None,
            "predictions":     None,
            "n_selected":      0,
            "n_exact_matches": 0,
            "status":          status,
        }
    
# ------------------------------------------------------------
    # Metric computation
    # ------------------------------------------------------------
    def compute_metric(self, neighborhood_result, z0, metric="mse",
                        tau_min=0.01, tau_max=0.5, n_thresholds=20):
        """
        Compute a per-feature stability metric from a neighborhood result.

        Parameters
        ----------
        neighborhood_result : dict
            Output of build_neighborhood for a single feature.
        z0 : float
            Reference prediction for the instance being explained.
        metric : {'mse', 'auc'}
            - 'mse': mean squared difference between neighbor predictions and z0
                        (lower = more stabilizing).
            - 'auc': area under the relative-proximity curve across tau thresholds
                        (higher = more stabilizing).
        tau_min, tau_max, n_thresholds : floats/int
            Used only when metric='auc'.

        Returns
        -------
        float or None
            Metric value, or None if the neighborhood is empty / invalid.
        """
        if neighborhood_result["status"] != "ok" or neighborhood_result["predictions"] is None:
            return None

        z = neighborhood_result["predictions"]
        z0 = float(z0)

        if metric == "mse":
            return float(np.mean((z - z0) ** 2))

        elif metric == "auc":
            denom = max(abs(z0), 1e-10)
            rel_diff = np.abs(z - z0) / denom
            taus = np.linspace(tau_min, tau_max, n_thresholds)
            prox = (rel_diff[:, None] <= taus[None, :]).mean(axis=0)  # (n_thresholds,)
            return float(np.trapezoid(prox, x=taus))

        else:
            raise ValueError(f"Unknown metric: {metric}. Use 'mse' or 'auc'.")
        

    # ------------------------------------------------------------
    # Convenience: neighborhoods + metrics for a single instance
    # ------------------------------------------------------------
    def explain_instance(self, x0, n_max=1000, metric="mse",
                         tau_min=0.01, tau_max=0.5, n_thresholds=20, rng=None):
        """
        Build neighborhoods and compute the chosen metric for every feature
        of a single instance.

        Parameters
        ----------
        x0 : array-like
            The instance to explain (1D, length = n_columns).
        n_max : int
            Target neighborhood size.
        metric : {'mse', 'auc'}
            Metric to compute per feature.
        tau_min, tau_max, n_thresholds :
            Passed to compute_metric (used only for 'auc').
        rng : np.random.Generator, optional
            Shared RNG for reproducibility across features.

        Returns
        -------
        dict
            {
                'z0':            float,
                'neighborhoods': {feature_name: neighborhood_result_dict},
                'metrics':       {feature_name: float or None},
            }
        """
        x0 = np.asarray(x0, dtype=np.float32).reshape(-1)
        z0 = float(self.predict(x0)[0])

        neighborhoods = {}
        metrics = {}
        for spec in self.feature_specs:
            fname = spec["name"]
            r = self.build_neighborhood(x0, fname, n_max=n_max, rng=rng)
            neighborhoods[fname] = r
            metrics[fname] = self.compute_metric(
                r, z0,
                metric=metric,
                tau_min=tau_min,
                tau_max=tau_max,
                n_thresholds=n_thresholds,
            )

        return {
            "z0":            z0,
            "neighborhoods": neighborhoods,
            "metrics":       metrics,
        }
    


import pandas as pd
import numpy as np
import joblib
 
# Required so joblib can resolve the custom function stored in the pipeline
from utils.preprocessing_utils import log1p_base10  # noqa: F401
 
 
def load_processed_data(
    attack,
    benign_data_path="Data/benign_only/all_days_benign.csv",
    attack_data_dir="Data/attacks_only",
    preprocessing_dir="_prepcosessing_artefacts/",
    model_dir_template="checkpoints_MLP/{attack}/",
    seed=42,
):
    model_dir = model_dir_template.format(attack=attack)

    # Load preprocessing artifacts
    ohe              = joblib.load(preprocessing_dir + "onehot_encoder.pkl")
    schema           = joblib.load(preprocessing_dir + "column_schema.pkl")
    numeric_pipeline = joblib.load(model_dir         + "numeric_pipeline.pkl")

    categorical_cols  = schema["categorical_cols"]
    numerical_cols    = schema["numerical_cols"]
    ohe_feature_names = schema["ohe_feature_names"]

    # Read raw CSVs
    df_benign = pd.read_csv(benign_data_path)
    df_attack = pd.read_csv(f"{attack_data_dir}/{attack}.csv")

    # One-hot encode categoricals
    benign_cat = pd.DataFrame(
        ohe.transform(df_benign[categorical_cols]),
        columns=ohe_feature_names,
        index=df_benign.index,
    )
    attack_cat = pd.DataFrame(
        ohe.transform(df_attack[categorical_cols]),
        columns=ohe_feature_names,
        index=df_attack.index,
    )

    # Concatenate numerical + categorical
    df_benign_proc = pd.concat(
        [df_benign[numerical_cols].reset_index(drop=True),
         benign_cat.reset_index(drop=True)],
        axis=1,
    )
    df_attack_proc = pd.concat(
        [df_attack[numerical_cols].reset_index(drop=True),
         attack_cat.reset_index(drop=True)],
        axis=1,
    )

    # Filter invalid rows (negatives or non-finite values in numerical columns)
    df_benign_proc = df_benign_proc[
        (df_benign_proc[numerical_cols] >= 0).all(axis=1)
        & np.isfinite(df_benign_proc[numerical_cols]).all(axis=1)
    ].copy()
    df_attack_proc = df_attack_proc[
        (df_attack_proc[numerical_cols] >= 0).all(axis=1)
        & np.isfinite(df_attack_proc[numerical_cols]).all(axis=1)
    ].copy()

    # Add labels and merge
    df_benign_proc["label"] = 0
    df_attack_proc["label"] = 1
    df_full = pd.concat([df_benign_proc, df_attack_proc], axis=0, ignore_index=True)
    df_full = df_full.sample(frac=1, random_state=seed).reset_index(drop=True)

    # Split features and labels
    X = df_full.drop(columns=["label"])
    y = df_full["label"]

    # Apply numeric pipeline (log1p + min-max)
    X_num = pd.DataFrame(
        numeric_pipeline.transform(X[numerical_cols]),
        columns=numerical_cols,
        index=X.index,
    )
    X_final = pd.concat([X_num, X[ohe_feature_names]], axis=1)

    return {
        "X_final": X_final,
        "y":       y,
        "schema":  schema,
    }