import pandas as pd
import pickle
from tensorflow.keras.models import load_model
from Feature import extract_url_features
from collections import Counter
from scipy.stats import entropy
import tensorflow as tf

# 🔹 URL 엔트로피 계산 함수
def calculate_url_entropy(url):
    counter = Counter(url)
    probabilities = [count / len(url) for count in counter.values()]
    return entropy(probabilities, base=2)

# 🔹 스케일러 불러오기
with open("scaler.pkl", "rb") as f:
    scaler = pickle.load(f)

# 🔹 모델 불러오기
model = load_model("best_model.h5")

# 🔹 예측 함수
@tf.function(reduce_retracing=True)
def predict_with_model(model, input_data):
    return model(input_data)

# 🔹 입력 URL 받기
url = input("URL입력 : ")

# 🔹 Feature.py에서 피처 추출
features = extract_url_features(url)

# 🔹 누락된 피처 보완
features['url_length'] = len(url)
features['url_entropy'] = calculate_url_entropy(url)

# 🔹 데이터프레임 생성 및 정렬
input_df = pd.DataFrame([features])
expected_columns = list(scaler.feature_names_in_)
input_df = input_df[expected_columns]

# 🔹 스케일링
input_scaled = scaler.transform(input_df)

# 🔹 예측
prediction = predict_with_model(model, input_scaled)
score = float(prediction.numpy()[0][0])  # 🔥 정확히 float으로 변환

# 🔹 출력
threshold = 0.5
if score > threshold:
    print(f"악성 (악성일 확률: {score:.4f})")
else:
    print(f"정상 (정상일 확률: {1 - score:.4f})")