model/model.py

"""
Model prediction helper module for phishing URL detection.
Handles model loading, feature extraction, and prediction.
"""

import os
import sys
import numpy as np
import pandas as pd
import joblib
import warnings
warnings.filterwarnings("ignore", message="X does not have valid feature names", category=UserWarning)

# Add parent directory to path to import url_feature_extraction module
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
from url_feature_extraction.url_feature_extractor import extract_features

# Global variable to cache the loaded model
_model_cache = None


def load_model(model_path="model/url_stacking_model.joblib"):
    """
    Load the saved stacking model from file.

    Args:
        model_path (str): Path to the model file relative to the FastAPI app directory

    Returns:
        dict: Dictionary containing model components:
            - base_models: Dictionary of base models
            - meta_scaler: Scaler for meta features
            - meta_model: Meta model for final prediction
            - feature_names: List of feature names
            - model_names: List of model names
    """
    global _model_cache

    # Return cached model if already loaded
    if _model_cache is not None:
        return _model_cache

    # Construct absolute path to model file
    current_dir = os.path.dirname(os.path.abspath(__file__))
    full_model_path = os.path.join(current_dir, "..", model_path)
    full_model_path = os.path.normpath(full_model_path)

    if not os.path.exists(full_model_path):
        raise FileNotFoundError(f"Model file not found at: {full_model_path}")

    # Load model
    model_data = joblib.load(full_model_path)
    print(f"✅ Model loaded successfully from: {full_model_path}")

    _model_cache = {
        "base_models": model_data["base_models"],
        "meta_scaler": model_data["meta_scaler"],
        "meta_model": model_data["meta_model"],
        "feature_names": model_data["feature_names"],
        "model_names": model_data["model_names"]
    }

    return _model_cache


def predict_url(url: str, model_components: dict = None):
    """
    Make prediction for a given URL.

    This function:
    1. Extracts features from the raw URL using url_feature_extractor
    2. Converts features to the format expected by the model
    3. Makes prediction using the stacking model

    Args:
        url (str): Raw URL to predict
        model_components (dict, optional): Pre-loaded model components.
                                          If None, will load the model.

    Returns:
        dict: Dictionary containing:
            - url: The input URL
            - predicted_label: 0 (legitimate) or 1 (phishing)
            - prediction: "legitimate" or "phishing"
            - phish_probability: Probability of being phishing (0.0 to 1.0)
            - confidence: Confidence percentage
            - features_extracted: Boolean indicating if features were successfully extracted
    """
    # Load model if not provided
    if model_components is None:
        model_components = load_model()

    # Extract features from URL
    features_dict = extract_features(url)

    # Check if feature extraction was successful
    if features_dict.get('has_title') is None:
        # URL was unreachable or feature extraction failed
        return {
            "url": url,
            "predicted_label": None,
            "prediction": "unknown",
            "phish_probability": None,
            "confidence": None,
            "features_extracted": False,
            "error": "Failed to extract features from URL. The URL may be unreachable or invalid."
        }

    # Make prediction using the features
    try:
        prediction_result = predict_from_features(features_dict, model_components)

        predicted_label = prediction_result["predicted_label"]
        phish_probability = prediction_result["phish_probability"]

        # Calculate confidence
        confidence = max(phish_probability, 1 - phish_probability) * 100

        return {
            "url": url,
            "predicted_label": predicted_label,
            "prediction": "phishing" if predicted_label == 1 else "legitimate",
            "phish_probability": round(phish_probability, 4),
            "confidence": round(confidence, 2),
            "features_extracted": True
        }
    except Exception as e:
        return {
            "url": url,
            "predicted_label": None,
            "prediction": "error",
            "phish_probability": None,
            "confidence": None,
            "features_extracted": True,
            "error": f"Prediction error: {str(e)}"
        }


def predict_from_features(features_dict: dict, model_components: dict):
    """
    Make predictions given a dictionary of extracted features.

    This function implements the stacking model prediction:
    - Level 0: Base models make predictions
    - Level 1: Meta model combines base model predictions

    Args:
        features_dict (dict): Dictionary where keys are feature names and values are feature values
        model_components (dict): The loaded components returned by load_model()

    Returns:
        dict: Contains 'predicted_label' (0 or 1) and 'phish_probability' (float)
    """
    base_models = model_components["base_models"]
    meta_scaler = model_components["meta_scaler"]
    meta_model = model_components["meta_model"]
    feature_names = model_components["feature_names"]
    model_names = model_components["model_names"]

    # Convert to DataFrame to ensure shape consistency
    X = pd.DataFrame([features_dict])

    # Ensure all required columns exist
    missing_cols = set(feature_names) - set(X.columns)
    if missing_cols:
        raise ValueError(f"❌ Missing required features: {missing_cols}")

    # Keep only known features and order them correctly
    X = X[feature_names]

    # ------------------------------
    # Level 0: Base model predictions
    # ------------------------------
    meta_features = np.zeros((X.shape[0], len(base_models)))
    for idx, (model_name, model) in enumerate(base_models.items()):
        meta_features[:, idx] = model.predict_proba(X)[:, 1]

    meta_features_df = pd.DataFrame(meta_features, columns=[f"{n}_pred" for n in model_names])

    # ------------------------------
    # Level 1: Meta-model prediction
    # ------------------------------
    meta_scaled = meta_scaler.transform(meta_features_df)
    meta_scaled = pd.DataFrame(meta_scaled, columns=meta_features_df.columns)

    final_pred = meta_model.predict(meta_scaled)[0]
    final_prob = meta_model.predict_proba(meta_scaled)[:, 1][0]

    return {
        "predicted_label": int(final_pred),
        "phish_probability": float(final_prob)
    }