import streamlit as st
from transformers import pipeline
from bert_score import score
import torch
import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
from rouge import Rouge

# Load the model once at app startup
@st.cache_resource
def load_embedding_model():
    return SentenceTransformer('all-MiniLM-L6-v2')

embedding_model = load_embedding_model()
rouge = Rouge()
# 还有我们自己的模型
# 在fine-tuned model上传到huggingface后，可以用pipeline来加载模型
# pipeline = pipeline(model="your_username/your_model_name")
# 之后我们可以继续用pipeline来直接调用模型得到prediction, 详见后

# Just using F1 for simplicity now
def calculate_bert_score(candidate, reference):
    # Make sure inputs are in list format as required by BERTScore
    if isinstance(candidate, str):
        candidate = [candidate]
    if isinstance(reference, str):
        reference = [reference]
        
    # Calculate the score
    P, R, F1 = score(candidate, reference, lang="en", return_hash=False)
    # Return F1 as a float (converting from tensor)
    return F1.item()

# Function to calculate cosine similarity
def calculate_cosine_similarity(text1, text2):
    # Get embeddings
    embedding1 = embedding_model.encode([text1])[0]
    embedding2 = embedding_model.encode([text2])[0]
    
    # Reshape for sklearn's cosine_similarity
    embedding1 = embedding1.reshape(1, -1)
    embedding2 = embedding2.reshape(1, -1)
    
    # Calculate similarity
    similarity = cosine_similarity(embedding1, embedding2)[0][0]
    return similarity

# Function to calculate ROUGE scores
def calculate_rouge_scores(candidate, reference):
    try:
        scores = rouge.get_scores(candidate, reference)[0]
        return {
            'rouge-1': scores['rouge-1']['f'],
            'rouge-l': scores['rouge-l']['f']
        }
    except:
        # Handle empty strings or other issues
        return {'rouge-1': 0, 'rouge-l': 0}

# UI
st.title("ResumeTailor")
st.write("Optimize your resume for specific job descriptions")

col1, col2 = st.columns(2)

with col1:
    st.header("Inputs")
    resume_text = st.text_area(
        "Paste your resume work experience here",
        height=300,
        placeholder="Describe your work experience..."
    )
    
    job_description = st.text_area(
        "Paste the job description here",
        height=300,
        placeholder="Paste the job description..."
    )
    
    if st.button("Refine Resume", type="primary"):
        # 在这里调用模型
        # prediction = pipeline(resume_text, job_description)
        refined_resume = "This would be the refined resume after model processing" # 现在这个是hard code的prediction
        refined_score = 0.85  # Example hardcoded score
        # 其实最终应该是下面的算法
        # original_bertscore = calculate_bert_score(resume_text, job_description)
        # refined_bertscore = calculate_bert_score(refined_resume, job_description)
        # delta_bertscore = refined_score - original_score

        # original_cosine_sim = calculate_cosine_similarity(resume_text, job_description)
        # refined_cosine_sim = calculate_cosine_similarity(refined_resume, job_description)
        # delta_cosine_sim = refined_cosine_sim - original_cosine_sim

        # original_rouge = calculate_rouge_scores(resume_text, job_description)
        # refined_rouge = calculate_rouge_scores(refined_resume, job_description)
        # delta_rouge1 = refined_rouge['rouge-1'] - original_rouge['rouge-1']
        # delta_rougel = refined_rouge['rouge-l'] - original_rouge['rouge-l']
        
        with col2:
            st.header("Results")
            st.text_area("Refined Work Experience", value=refined_resume, height=300)
            st.subheader("Similarity Score")
            # Below are just hardcoded for now
            # Use 2 x 2 matrix to show the scores
            a, b = st.columns(2)
            with a:
                st.metric("BERT Score (F1)", value=f"{refined_score:.2f}", delta=f"{refined_score - 0.85:.2f}")
                st.metric("Cosine Similarity", value=f"{0.90:.2f}", delta=f"{0.90 - 0.85:.2f}")
            with b:
                st.metric("ROUGE-1 (F1)", value=f"{0.88:.2f}", delta=f"{0.88 - 0.85:.2f}")
                st.metric("ROUGE-L (F1)", value=f"{0.87:.2f}", delta=f"{0.87 - 0.85:.2f}")
    
            # 其实应该是
            # st.metric("BERT Score (F1)", value=f"{refined_bertscore:.2f}", delta=f"{delta_score:.2f}")
            # st.metric("Cosine Similarity", value=f"{refined_cosine_sim:.2f}", delta=f"{delta_cosine_sim:.2f}")
            # st.metric("ROUGE-1 (F1)", value=f"{refined_rouge['rouge-1']:.2f}", delta=f"{delta_rouge1:.2f}")
            # st.metric("ROUGE-L (F1)", value=f"{refined_rouge['rouge-l']:.2f}", delta=f"{delta_rougel:.2f}")






# Below is a demo code to use hot-dog classification model

# from PIL import Image

# pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")

# st.title("Hot Dog? Or Not?")

# file_name = st.file_uploader("Upload a hot dog candidate image")

# if file_name is not None:
#     col1, col2 = st.columns(2)

#     image = Image.open(file_name)
#     col1.image(image, use_column_width=True)
#     predictions = pipeline(image)

#     col2.header("Probabilities")
#     for p in predictions:
#         col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")