inference.py

"""Real-time Inference Engine for Neuromorphic EDR Sentinel"""
import torch
import numpy as np
import time
import json
import os
from typing import Dict, List, Optional, Tuple
from collections import deque
from dataclasses import dataclass


@dataclass
class DetectionResult:
    """Detection result from inference"""
    timestamp: float
    is_anomalous: bool
    confidence: float
    anomaly_score: float
    processing_time_ms: float
    sequence_id: str
    details: Dict


class InferenceEngine:
    """Real-time inference engine for EDR detection"""
    
    def __init__(self, model_path: str, config: Dict = None):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Load model
        checkpoint = torch.load(model_path, map_location=self.device)
        self.config = checkpoint.get("config", config or {})
        
        # Initialize model
        from modules.kan.kolmogorov_arnold import SNN_KAN_Hybrid
        self.model = SNN_KAN_Hybrid(
            input_dim=self.config.get("feature_dim", 128),
            snn_hidden=self.config.get("snn_hidden", [128, 64]),
            kan_hidden=self.config.get("kan_hidden", [64, 32]),
            num_classes=2,
            threshold=self.config.get("threshold", 0.5),
            membrane_decay=self.config.get("membrane_decay", 0.9),
        ).to(self.device)
        
        self.model.load_state_dict(checkpoint["model_state_dict"])
        self.model.eval()
        
        # Sequence buffer for streaming detection
        self.sequence_buffer = deque(maxlen=64)
        self.detection_threshold = 0.5
        
        # Performance tracking
        self.inference_times = deque(maxlen=1000)
        self.total_detections = 0
        self.total_anomalies = 0
        
        print(f"Model loaded from {model_path}")
        print(f"Device: {self.device}")
        print(f"Parameters: {sum(p.numel() for p in self.model.parameters()):,}")
        
    def preprocess_sequence(self, syscalls: List[Dict]) -> np.ndarray:
        """Preprocess syscall sequence to feature vector"""
        features = np.zeros(128)
        
        if not syscalls:
            return features
            
        # Syscall frequency histogram
        freq_hist = np.zeros(150)
        for sc in syscalls:
            sc_id = sc.get("syscall_id", 0)
            if sc_id < 150:
                freq_hist[sc_id] += 1
        freq_hist = freq_hist / (freq_hist.sum() + 1e-8)
        features[:150] = freq_hist[:128]
        
        # Temporal features
        timestamps = [sc.get("timestamp", 0) for sc in syscalls]
        if len(timestamps) > 1:
            intervals = np.diff(timestamps)
            features[120:128] = [
                np.mean(intervals), np.std(intervals),
                np.min(intervals), np.max(intervals),
                np.median(intervals),
                np.percentile(intervals, 25),
                np.percentile(intervals, 75),
                len(syscalls) / 100.0,
            ][:8]
            
        return features.astype(np.float32)
    
    def detect(self, syscalls: List[Dict], sequence_id: str = "") -> DetectionResult:
        """Run detection on a syscall sequence"""
        start_time = time.time()
        
        # Preprocess
        features = self.preprocess_sequence(syscalls)
        features_tensor = torch.FloatTensor(features).unsqueeze(0).to(self.device)
        
        # Inference
        with torch.no_grad():
            output, snn_logits, kan_logits = self.model(features_tensor)
            probability = torch.sigmoid(output).item()
            
        processing_time = (time.time() - start_time) * 1000  # ms
        self.inference_times.append(processing_time)
        
        is_anomalous = probability > self.detection_threshold
        confidence = abs(probability - 0.5) * 2.0  # Normalize to [0, 1]
        
        self.total_detections += 1
        if is_anomalous:
            self.total_anomalies += 1
            
        return DetectionResult(
            timestamp=time.time(),
            is_anomalous=is_anomalous,
            confidence=confidence,
            anomaly_score=probability,
            processing_time_ms=processing_time,
            sequence_id=sequence_id,
            details={
                "snn_score": torch.sigmoid(snn_logits).cpu().numpy().tolist(),
                "kan_score": torch.sigmoid(kan_logits).cpu().numpy().tolist(),
                "num_syscalls": len(syscalls),
            }
        )
    
    def detect_batch(self, sequences: List[List[Dict]]) -> List[DetectionResult]:
        """Batch detection for multiple sequences"""
        features_list = []
        for seq in sequences:
            features = self.preprocess_sequence(seq)
            features_list.append(features)
            
        features_tensor = torch.FloatTensor(features_list).to(self.device)
        
        with torch.no_grad():
            output, _, _ = self.model(features_tensor)
            probabilities = torch.sigmoid(output).cpu().numpy()
            
        results = []
        for i, prob in enumerate(probabilities):
            is_anomalous = prob > self.detection_threshold
            results.append(DetectionResult(
                timestamp=time.time(),
                is_anomalous=is_anomalous,
                confidence=abs(prob - 0.5) * 2.0,
                anomaly_score=prob,
                processing_time_ms=0,
                sequence_id=f"batch_{i}",
                details={}
            ))
            
        return results
    
    def get_performance_stats(self) -> Dict:
        """Get inference performance statistics"""
        if not self.inference_times:
            return {"avg_latency_ms": 0, "p99_latency_ms": 0, "total_detections": 0}
            
        times = np.array(self.inference_times)
        return {
            "avg_latency_ms": float(np.mean(times)),
            "p50_latency_ms": float(np.percentile(times, 50)),
            "p95_latency_ms": float(np.percentile(times, 95)),
            "p99_latency_ms": float(np.percentile(times, 99)),
            "total_detections": self.total_detections,
            "total_anomalies": self.total_anomalies,
            "anomaly_rate": self.total_anomalies / (self.total_detections + 1e-8),
        }
    
    def export_model_for_deployment(self, output_path: str):
        """Export model for deployment (quantized)"""
        # Quantize model
        self.model.cpu()
        quantized_model = torch.quantization.quantize_dynamic(
            self.model, {torch.nn.Linear}, dtype=torch.qint8
        )
        
        torch.save({
            "model_state_dict": quantized_model.state_dict(),
            "config": self.config,
            "quantized": True,
        }, output_path)
        
        print(f"Quantized model exported to {output_path}")
        
        # Get size
        size_mb = os.path.getsize(output_path) / (1024 * 1024)
        print(f"Model size: {size_mb:.2f} MB")