import streamlit as st
from sentence_transformers import SentenceTransformer
import joblib
import numpy as np

# ---------------------------------------------------
# Load SBERT + Classifier + LabelEncoder
# ---------------------------------------------------

@st.cache_resource
def load_models():
    # Must match training model exactly
    embedder = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

    # Load classifier
    classifier = joblib.load("classifier.pkl")

    # Load label encoder
    label_encoder = joblib.load("label_encoder.pkl")

    return embedder, classifier, label_encoder


embedder, classifier, label_encoder = load_models()

# ---------------------------------------------------
# Streamlit UI
# ---------------------------------------------------

st.set_page_config(page_title="Trademark Class Predictor", page_icon="🔍")

st.title("🔍 NICE Class Predictor (3, 5, Both)")

description = st.text_area("Enter product description:", height=150)

if st.button("Predict Class"):
    if description.strip() == "":
        st.warning("⚠️ Please enter a valid description.")
    else:
        # Embed the input
        emb = embedder.encode([description])

        # Predict (numeric)
        pred_num = classifier.predict(emb)[0]

        # Convert to readable class label
        pred_label = label_encoder.inverse_transform([pred_num])[0]

        # Show result
        st.subheader("Prediction:")
        if pred_label == "3_only":
            st.success("🧴 Class 3 — Cosmetics & Cleaning Preparations")
        elif pred_label == "5_only":
            st.success("💊 Class 5 — Pharmaceutical & Medical Products")
        elif pred_label == "both":
            st.success("🔄 Both — Mixed Class 3 + Class 5")
        else:
            st.error("Unknown label.")

        # Confidence scores
        if hasattr(classifier, "predict_proba"):
            proba = classifier.predict_proba(emb)[0]
            st.write("### Confidence Scores:")
            for cls, p in zip(label_encoder.classes_, proba):
                st.write(f"- **{cls}**: `{p:.3f}`")