import streamlit as st
import joblib
import numpy as np
import string

# Load the pipeline and label mapping
pipeline = joblib.load("ensemble_pipeline.joblib")
label_map = joblib.load("label_map.joblib")
# Create a reverse mapping for nice output
reverse_label_map = {v: k for k, v in label_map.items()}

# App title and description
st.title("Password Strength Predictor")
st.write("""
Enter a password below to see its predicted strength (Weak, Medium, or Strong)
using a pre-trained ensemble classifier.
""")

# Text input for the password
password = st.text_input("Enter a Password:")

# Optional: Define the same numerical feature function if needed
def generate_numerical_features(pwd):
    return np.array([
        len(pwd),
        sum(1 for char in pwd if char.islower()) / max(1, len(pwd)),
        sum(1 for char in pwd if char.isupper()) / max(1, len(pwd)),
        sum(1 for char in pwd if char.isdigit()) / max(1, len(pwd)),
        sum(1 for char in pwd if not char.isalnum()) / max(1, len(pwd)),
        int(any(char in string.punctuation for char in pwd))
    ])

# Prediction action
if st.button("Predict Strength"):
    if password:
        # Use the pipeline to predict directly from the raw password
        pred_numeric = pipeline.predict([password])
        # Convert numeric output to a human-readable label using the reverse mapping
        pred_label = reverse_label_map.get(pred_numeric[0], "Unknown")
        st.success(f"Predicted Password Strength: **{pred_label}**")
    else:
        st.warning("Please enter a password to get a prediction.")