"""
preprocess.py  -  TemporalDrift-ETM
=====================================
Converts any NTLFlowLyzer-compatible CSV into the fixed 35-feature
scaled input that the ensemble model expects.

The model was trained on 35 features selected by Random Forest importance
ranking from an original 76-column feature space.  Uploaded CSV files may
have any number of columns (40, 50, 70, 85, ...).  This module handles:

    - Any number of input columns
    - Columns that are metadata, identifiers, or labels  -> silently skipped
    - Selected features present in the CSV               -> cleaned and scaled
    - Selected features absent from the CSV              -> filled with 0.0
    - Non-numeric cell values                            -> coerced to column mean
    - Inf / -Inf values                                  -> replaced with column mean

The per-feature scaling parameters (data_min, scale) are extracted from
the training MinMaxScaler at notebook-save time and stored in scaler_35.pkl,
so this module never requires the full 76-feature scaler at inference time.

Standalone usage
----------------
    python preprocess.py input.csv output.csv [--models models/]

Module usage (from app.py)
--------------------------
    from preprocess import load_scaler_35, prepare_features
    scaler_35        = load_scaler_35("models/scaler_35.pkl")
    X, present, miss = prepare_features(df, scaler_35)
"""

import os
import sys
import numpy as np
import pandas as pd
import joblib


# ---------------------------------------------------------------------------
# Column names that are never network-flow features.
# Any CSV column in this set is silently ignored during feature extraction.
# ---------------------------------------------------------------------------

NON_FEATURE_COLS = frozenset({
    "Flow ID",
    "Src IP",
    "Src Port",
    "Dst IP",
    "Dst Port",
    "Protocol",
    "Timestamp",
    "Malware Family",
    "Label",
    "label",
    "session",
    "Session",
    "Prediction",
})


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def load_scaler_35(path):
    """
    Load the 35-feature scaler dict produced by the notebook save cell.

    The dict contains three keys:

        feature_names : list[str]   - the 35 selected feature names (in order)
        data_min      : np.ndarray  - training minimum per feature, shape (35,)
        scale         : np.ndarray  - 1/(max-min) per feature,      shape (35,)

    Parameters
    ----------
    path : str
        Path to scaler_35.pkl.

    Returns
    -------
    dict
    """
    artifact = joblib.load(path)
    required = {"feature_names", "data_min", "scale"}
    missing  = required - set(artifact.keys())
    if missing:
        raise ValueError(
            f"scaler_35.pkl is missing keys: {missing}. "
            "Re-run the updated notebook save cell to regenerate it."
        )
    return artifact


def prepare_features(df, scaler_35):
    """
    Convert a DataFrame with any columns into the 35-feature scaled matrix.

    For each of the 35 model features:
      - Column present in df  -> coerce to float, replace non-finite values
                                 with the column mean, apply MinMaxScaler
                                 formula: (x - data_min) * scale, clip to [0,1].
      - Column absent from df -> fill the entire column with 0.0 in scaled
                                 space (neutral value that does not push the
                                 model toward any particular prediction).

    Parameters
    ----------
    df        : pd.DataFrame
        Raw input data (any number of columns, any column names).
    scaler_35 : dict
        Output of load_scaler_35().

    Returns
    -------
    X_scaled : np.ndarray, shape (n_samples, 35)
        Model-ready scaled feature matrix.
    present  : list[str]
        Names of the 35 features that were found in df.
    missing  : list[str]
        Names of the 35 features that were absent in df (filled with 0.0).
    """
    feat_names = scaler_35["feature_names"]
    data_min   = np.asarray(scaler_35["data_min"], dtype=np.float64)
    scale      = np.asarray(scaler_35["scale"],    dtype=np.float64)

    n       = len(df)
    X       = np.zeros((n, len(feat_names)), dtype=np.float64)
    present = []
    missing = []

    for i, feat in enumerate(feat_names):
        if feat not in df.columns:
            missing.append(feat)
            # Column stays at 0.0 - neutral fill in scaled space.
            continue

        col = pd.to_numeric(df[feat], errors="coerce")

        # Replace non-finite values with the finite column mean.
        finite_vals = col[np.isfinite(col)]
        col_mean    = float(finite_vals.mean()) if len(finite_vals) > 0 else 0.0
        col         = col.fillna(col_mean)
        col         = col.replace([np.inf, -np.inf], col_mean)

        # MinMaxScaler transform: X_scaled = (X - data_min) * scale
        # This exactly replicates sklearn's MinMaxScaler.transform()
        # but for a single feature at a time.
        X[:, i] = (col.to_numpy(dtype=np.float64) - data_min[i]) * scale[i]
        present.append(feat)

    # Clip to [0, 1] as sklearn's MinMaxScaler.transform() does.
    np.clip(X, 0.0, 1.0, out=X)

    return X, present, missing


def report_coverage(present, missing, file=None):
    """
    Print a human-readable feature coverage summary.

    Parameters
    ----------
    present : list[str]
    missing : list[str]
    file    : file object or None (default: stdout)
    """
    total = len(present) + len(missing)
    print(f"Feature coverage: {len(present)}/{total} columns found in CSV.", file=file)
    if missing:
        print(f"  {len(missing)} feature(s) absent - filled with 0 in scaled space:",
              file=file)
        for m in missing:
            print(f"    - {m}", file=file)


# ---------------------------------------------------------------------------
# Standalone CLI
# ---------------------------------------------------------------------------

def _main():
    import argparse

    parser = argparse.ArgumentParser(
        description=(
            "Convert any NTLFlowLyzer-compatible CSV to the 35-feature "
            "scaled format expected by TemporalDrift-ETM."
        )
    )
    parser.add_argument("input",  help="Path to the raw input CSV file")
    parser.add_argument("output", help="Path to write the processed output CSV")
    parser.add_argument(
        "--models", default="models",
        help="Directory containing scaler_35.pkl (default: models/)",
    )
    args = parser.parse_args()

    scaler_35_path = os.path.join(args.models, "scaler_35.pkl")
    if not os.path.exists(scaler_35_path):
        print(
            f"ERROR: {scaler_35_path} not found.\n"
            "Re-run the updated notebook save cell to generate scaler_35.pkl.",
            file=sys.stderr,
        )
        sys.exit(1)

    print(f"Reading {args.input} ...")
    df = pd.read_csv(args.input)
    print(f"  {len(df):,} rows  |  {df.shape[1]} columns")

    scaler_35  = load_scaler_35(scaler_35_path)
    feat_names = scaler_35["feature_names"]
    print(f"\nPreparing {len(feat_names)} model features ...")

    X, present, missing = prepare_features(df, scaler_35)
    report_coverage(present, missing)

    out_df = pd.DataFrame(X, columns=feat_names)
    out_df.to_csv(args.output, index=False)

    print(f"\nSaved: {args.output}")
    print(f"Output shape: {out_df.shape}")


if __name__ == "__main__":
    _main()