from transformers import Pipeline
import numpy as np
import joblib
from typing import Dict, List, Union

class FraudDetectionPipeline(Pipeline):
    def __init__(self):
        super().__init__()
        # Load the model and scaler
        self.model = joblib.load("random_forest_model.joblib")
        self.scaler = joblib.load("rf_scaler.joblib")
    
    def preprocess(self, features: Dict[str, Union[int, float]]) -> np.ndarray:
        """
        Preprocess the input features
        Expected features:
        - account_age: int (months)
        - cred_changes_freq: float (per year)
        - return_order_ratio: float
        - vpn_usage: int (0 or 1)
        - credit_score: int
        """
        # Convert input to correct format
        input_data = np.array([[
            features['account_age'],
            features['cred_changes_freq'],
            features['return_order_ratio'],
            features['vpn_usage'],
            features['credit_score']
        ]])
        
        # Scale the features
        scaled_input = self.scaler.transform(input_data)
        return scaled_input
    
    def _forward(self, features: Dict[str, Union[int, float]]) -> Dict[str, Union[str, float]]:
        """
        Make prediction using the model
        """
        # Preprocess
        scaled_input = self.preprocess(features)
        
        # Get prediction and probability
        prediction = self.model.predict(scaled_input)[0]
        probabilities = self.model.predict_proba(scaled_input)[0]
        
        # Return prediction and confidence
        return {
            "prediction": "Fraud" if prediction == 1 else "Not Fraud",
            "confidence": float(probabilities[prediction]),
            "fraud_probability": float(probabilities[1])
        }
    
    def postprocess(self, model_outputs):
        return model_outputs

def load_pipeline():
    return FraudDetectionPipeline()