import streamlit as st
import joblib
import pandas as pd
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

# Load the necessary models
model = joblib.load("resume_classification_model.pkl")
tokenizer = joblib.load("resume_label_encoder.pkl")
sbert_model = SentenceTransformer('all-MiniLM-L6-v2')

# Set up the Streamlit page configuration
st.set_page_config(page_title="Multi Resume Analyzer", layout="centered")
st.markdown("<h1 style='text-align: center; color: #4CAF50;'>📝 Multi Resume Analyzer</h1>", unsafe_allow_html=True)

# Input job description
jd = st.text_input("Enter the Job Description:")

# Initialize session state
if "text_blocks" not in st.session_state:
    st.session_state.text_blocks = []
if "submitted" not in st.session_state:
    st.session_state.submitted = False
if "message" not in st.session_state:
    st.session_state.message = ""
if "text_input_key" not in st.session_state:
    st.session_state.text_input_key = 0

# Show textarea and buttons if not submitted
if not st.session_state.submitted:
    st.markdown("### ✍️ Enter or Paste Your Resume Below:")
    text_input = st.text_area("Text Input", height=200, key=f"input_{st.session_state.text_input_key}")

    col1, col2 = st.columns(2)
    with col1:
        if st.button("Next"):
            if text_input.strip():
                st.session_state.text_blocks.append(text_input.strip())
                st.session_state.message = f"✅ Resume block {len(st.session_state.text_blocks)} saved!"
                st.session_state.text_input_key += 1
                st.rerun()
            else:
                st.warning("Please enter some text before clicking Next.")
    with col2:
        if st.button("Submit"):
            if text_input.strip():
                st.session_state.text_blocks.append(text_input.strip())
            st.session_state.submitted = True
            st.rerun()

# Show feedback message after "Next"
if st.session_state.message:
    st.success(st.session_state.message)

# After submission
if st.session_state.submitted and jd:
    st.markdown("## ✅ Submission Complete!")

    resumes = st.session_state.text_blocks

    # SBERT embeddings
    embeddings = sbert_model.encode(resumes)  # one vector per resume
    jd_embedding = sbert_model.encode([jd])   # job description embedding

    # Cosine similarities
    similarities = cosine_similarity(jd_embedding, embeddings)[0]

    # Predict categories for all resumes
    predicted_labels = model.predict(embeddings)
    predicted_categories = tokenizer.inverse_transform(predicted_labels)

    # Create DataFrame
    df = pd.DataFrame({
        "Resume": [f"Resume {i+1}" for i in range(len(resumes))],
        "Index": list(range(len(resumes))),
        "Score": similarities,
        "Predicted Category": predicted_categories
    })

    # Ranking
    df["Rank"] = df["Score"].rank(ascending=False, method='first').astype(int)
    df = df.sort_values(by="Rank")

    # Predicted category of top resume
    top_resume_embedding = embeddings[df.iloc[0]["Index"]].reshape(1, -1)
    top_pred = model.predict(top_resume_embedding)
    top_category = tokenizer.inverse_transform(top_pred)[0]
    st.success(f"🎯 Predicted Job Category for Best Match: **{top_category}**")

    # Display full results table
    st.markdown("### 📊 Resume Analysis Results:")
    st.dataframe(df[['Resume', 'Index', 'Score', 'Rank', 'Predicted Category']], width=700)