import numpy as np
import torch
import joblib
from model.model import VOIPClassifier, get_top_features

scaler = joblib.load("scaler.pkl")

def load_model(model_path, input_dim):
    model = VOIPClassifier(input_dim)
    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
    model.eval()
    return model

def predict(model, input_list, scaler):
    input_array = np.array(input_list).reshape(1, -1)
    input_scaled = scaler.transform(input_array)
    input_tensor = torch.FloatTensor(input_scaled)
    with torch.no_grad():
        output = model(input_tensor)
        probability = output.item()
        predicted_class = 1 if probability >= 0.5 else 0
    return {"probability": probability, "class": predicted_class}

def inference(input):
    top_features = get_top_features()
    model = load_model("VOIP_Classifier.pth", input_dim=len(top_features))
    result = predict(model, input, scaler)
    return result

if __name__ == "__main__":
    top_features = get_top_features()
    model = load_model("VOIP_Classifier.pth", input_dim=len(top_features))
    sample_input = [2, 4, 5]
    print(inference(sample_input))