import os
os.environ['HF_HOME'] = '/tmp'
import time
import streamlit as st
import pandas as pd
import io
import plotly.express as px
import zipfile
import json
import hashlib


from typing import Optional
from gliner import GLiNER
from comet_ml import Experiment

# --- Page Configuration and UI Elements ---
st.set_page_config(layout="wide", page_title="NER")
st.subheader("HR.ai", divider="green")
st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")


st.markdown(
    """
    <style>
    /* Main app background and text color */
    .stApp {
        background-color: #F5FFFA; /* Mint cream, a very light green */
        color: #000000; /* Black for the text */
    }
    /* Sidebar background color */
    .css-1d36184 {
        background-color: #B2F2B2; /* A pale green for the sidebar */
        secondary-background-color: #B2F2B2;
    }
    /* Expander background color and header */
    .streamlit-expanderContent, .streamlit-expanderHeader {
        background-color: #F5FFFA;
    }
    /* Text Area background and text color */
    .stTextArea textarea {
        background-color: #D4F4D4; /* A light, soft green */
        color: #000000; /* Black for text */
    }
    /* Text Input background and text color */
    .stTextInput input {
        background-color: #D4F4D4; /* Same as the text area for consistency */
        color: #000000;
    }
    /* Button background and text color */
    .stButton > button {
        background-color: #D4F4D4;
        color: #000000;
    }
    /* Warning box background and text color */
    .stAlert.st-warning {
        background-color: #C8F0C8; /* A light green for the warning box */
        color: #000000;
    }
    /* Success box background and text color */
    .stAlert.st-success {
        background-color: #C8F0C8; /* A light green for the success box */
        color: #000000;
    }
    /* Tab color when active */
    .stTabs [data-baseweb="tab-list"] button[aria-selected="true"] {
        background-color: #D4F4D4;
        color: #000000;
    }
    
    </style>
    """,
    unsafe_allow_html=True
)

expander = st.expander("**Important notes**")
expander.write(""" 
**How to Use the HR.ai web app:** 

1. Type or paste your text into the text area, then press Ctrl + Enter.
2. Click the 'Results' button to extract and tag entities in your text data. 

Results are presented in easy-to-read tables, visualized in an interactive tree map, pie chart and bar chart, and are available for download along with a Glossary of tags.      

**How to Use the Question-Answering feature:**

1. Type or paste your text into the text area, then press Ctrl + Enter. 
2. Click the 'Add Question' button to add your question to the Record of Questions. You can manage your questions by deleting them one by one.
3. Click the 'Extract Answers' button to extract the answer to your question.

Results are presented in an easy-to-read table, visualized in an interactive tree map, and is available for download.   

**Entities:** "Email", "Phone_number", "Street_address", "City", "Country", "Date_of_birth", "Marital_status", "Person", "Full_time", "Part_time", "Contract", "Terminated", "Retired", "Job_title", "Date", "Organization", "Role", "Performance_score", "Leave_of_absence", "Retirement_plan", "Bonus", "Stock_options", "Health_insurance", "Pay_rate", "Annual_salary", "Tax", "Deductions", "Interview_type", "Applicant", "Referral", "Job_board", "Recruiter", "Offer_letter", "Agreement", "Certification", "Skill"

**Usage Limits:** You can request results unlimited times for one (1) month.      

**Supported Languages:** English     

**Technical issues:** If your connection times out, please refresh the page or reopen the app's URL. 

For any errors or inquiries, please contact us at info@nlpblogs.com""")


with st.sidebar:
    st.write("Use the following code to embed the web app on your website. Feel free to adjust the width and height values to fit your page.")
    code = '''
    <iframe
    src="https://aiecosystem-hr-ai.hf.space"
    frameborder="0"
    width="850"
    height="450"
    ></iframe>
    '''
    st.code(code, language="html")
    st.text("")
    st.text("")
    st.divider()
    st.subheader("🚀 Ready to build your own AI Web App?", divider="green")
    st.link_button("AI Web App Builder", "https://nlpblogs.com/build-your-named-entity-recognition-app/", type="primary")

# --- Comet ML Setup ---
COMET_API_KEY = os.environ.get("COMET_API_KEY")
COMET_WORKSPACE = os.environ.get("COMET_WORKSPACE")
COMET_PROJECT_NAME = os.environ.get("COMET_PROJECT_NAME")
comet_initialized = bool(COMET_API_KEY and COMET_WORKSPACE and COMET_PROJECT_NAME)

if not comet_initialized:
    st.warning("Comet ML not initialized. Check environment variables.")

# --- Model Loading and Caching ---
@st.cache_resource
def load_gliner_model(model_name):
    """Initializes and caches the GLiNER model."""
    try:
        if model_name == "HR_AI":
            return GLiNER.from_pretrained("knowledgator/gliner-multitask-large-v0.5", nested_ner=True, num_gen_sequences=2, gen_constraints=labels)
        elif model_name == "InfoFinder":
            return GLiNER.from_pretrained("knowledgator/gliner-multitask-large-v0.5", device="cpu")
    except Exception as e:
        st.error(f"Error loading the GLiNER model: {e}")
        st.stop()

# --- HR_AI Model Labels and Mappings ---
labels = ["Email", "Phone_number", "Street_address", "City", "Country", "Date_of_birth", "Marital_status", "Person", "Full_time", "Part_time", "Contract", "Terminated", "Retired", "Job_title", "Date", "Organization", "Role", "Performance_score", "Leave_of_absence", "Retirement_plan", "Bonus", "Stock_options", "Health_insurance", "Pay_rate", "Annual_salary", "Tax", "Deductions", "Interview_type", "Applicant", "Referral", "Job_board", "Recruiter", "Offer_letter", "Agreement", "Certification", "Skill"]

category_mapping = {
    "Contact Information": ["Email", "Phone_number", "Street_address", "City", "Country"],
    "Personal Details": ["Date_of_birth", "Marital_status", "Person"],
    "Employment Status": ["Full_time", "Part_time", "Contract", "Terminated", "Retired"],
    "Employment Information": ["Job_title", "Date", "Organization", "Role"],
    "Performance": ["Performance_score"],
    "Attendance": ["Leave_of_absence"],
    "Benefits": ["Retirement_plan", "Bonus", "Stock_options", "Health_insurance"],
    "Compensation": ["Pay_rate", "Annual_salary"],
    "Deductions": ["Tax", "Deductions"],
    "Recruitment & Sourcing": ["Interview_type", "Applicant", "Referral", "Job_board", "Recruiter"],
    "Legal & Compliance": ["Offer_letter", "Agreement"],
    "Professional_Development": ["Certification", "Skill"]
}
reverse_category_mapping = {label: category for category, label_list in category_mapping.items() for label in label_list}

# --- InfoFinder Helpers ---
if 'user_labels' not in st.session_state:
    st.session_state.user_labels = []

def get_stable_color(label):
    hash_object = hashlib.sha1(label.encode('utf-8'))
    hex_dig = hash_object.hexdigest()
    return '#' + hex_dig[:6]

# --- Main App with Tabs ---
tab1, tab2 = st.tabs(["HR.ai", "Question-Answering"])

with tab1:
    
    
    # Load model for this tab
    model_hr = load_gliner_model("HR_AI")

    text = st.text_area("Type or paste your text below, and then press Ctrl + Enter", height=250, key='my_text_area_hr')

    def clear_text_hr():
        st.session_state['my_text_area_hr'] = ""

    st.button("Clear text", on_click=clear_text_hr, key="clear_hr")

    if st.button("Results"):
        start_time = time.time()
        if not text.strip():
            st.warning("Please enter some text to extract entities.")
        else:
            with st.spinner("Extracting entities...", show_time=True):
                entities = model_hr.predict_entities(text, labels)
                df = pd.DataFrame(entities)

                if not df.empty:
                    df['category'] = df['label'].map(reverse_category_mapping)
                    if comet_initialized:
                        experiment = Experiment(api_key=COMET_API_KEY, workspace=COMET_WORKSPACE, project_name=COMET_PROJECT_NAME)
                        experiment.log_parameter("input_text", text)
                        experiment.log_table("predicted_entities", df)

                    st.subheader("Grouped Entities by Category", divider="green")
                    category_names = sorted(list(category_mapping.keys()))
                    category_tabs_hr = st.tabs(category_names)
                    for i, category_name in enumerate(category_names):
                        with category_tabs_hr[i]:
                            df_category_filtered = df[df['category'] == category_name]
                            if not df_category_filtered.empty:
                                st.dataframe(df_category_filtered.drop(columns=['category']), use_container_width=True)
                            else:
                                st.info(f"No entities found for the '{category_name}' category.")

                    with st.expander("See Glossary of tags"):
                        st.write('''
                        - **text**: ['entity extracted from your text data']
                        - **score**: ['accuracy score; how accurately a tag has been assigned to a given entity']
                        - **label**: ['label (tag) assigned to a given extracted entity']
                        - **start**: ['index of the start of the corresponding entity']
                        - **end**: ['index of the end of the corresponding entity']
                        ''')
                    st.divider()
                    
                    st.subheader("Candidate Card", divider="green")
                    fig_treemap = px.treemap(df, path=[px.Constant("all"), 'category', 'label', 'text'], values='score', color='category')
                    fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25), paper_bgcolor='#F5FFFA', plot_bgcolor='#F5FFFA')
                    st.plotly_chart(fig_treemap)
                    
                    col1, col2 = st.columns(2)
                    with col1:
                        st.subheader("Pie chart", divider="green")
                        grouped_counts = df['category'].value_counts().reset_index()
                        grouped_counts.columns = ['category', 'count']
                        fig_pie = px.pie(grouped_counts, values='count', names='category', hover_data=['count'], labels={'count': 'count'}, title='Percentage of predicted categories')
                        fig_pie.update_traces(textposition='inside', textinfo='percent+label')
                        fig_pie.update_layout(paper_bgcolor='#F5FFFA', plot_bgcolor='#F5FFFA')
                        st.plotly_chart(fig_pie)
                    
                    with col2:
                        st.subheader("Bar chart", divider="green")
                        fig_bar = px.bar(grouped_counts, x="count", y="category", color="category", text_auto=True, title='Occurrences of predicted categories')
                        fig_bar.update_layout(paper_bgcolor='#F5FFFA', plot_bgcolor='#F5FFFA')
                        st.plotly_chart(fig_bar)
                    
                    st.subheader("Most Frequent Entities", divider="green")
                    word_counts = df['text'].value_counts().reset_index()
                    word_counts.columns = ['Entity', 'Count']
                    repeating_entities = word_counts[word_counts['Count'] > 1]
                    if not repeating_entities.empty:
                        st.dataframe(repeating_entities, use_container_width=True)
                        fig_repeating_bar = px.bar(repeating_entities, x='Entity', y='Count', color='Entity')
                        fig_repeating_bar.update_layout(xaxis={'categoryorder': 'total descending'}, paper_bgcolor='#F5FFFA', plot_bgcolor='#F5FFFA')
                        st.plotly_chart(fig_repeating_bar)
                    else:
                        st.warning("No entities were found that occur more than once.")
                    
                    st.divider()
                    
                    dfa = pd.DataFrame(data={'Column Name': ['text', 'label', 'score', 'start', 'end'], 'Description': ['entity extracted from your text data', 'label (tag) assigned to a given extracted entity', 'accuracy score; how accurately a tag has been assigned to a given entity', 'index of the start of the corresponding entity', 'index of the end of the corresponding entity']})
                    buf = io.BytesIO()
                    with zipfile.ZipFile(buf, "w") as myzip:
                        myzip.writestr("Summary of the results.csv", df.to_csv(index=False))
                        myzip.writestr("Glossary of tags.csv", dfa.to_csv(index=False))
                    
                    st.download_button(
                        label="Download results and glossary (zip)",
                        data=buf.getvalue(),
                        file_name="nlpblogs_results.zip",
                        mime="application/zip",
                    )
                    
                    if comet_initialized:
                        experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
                        experiment.end()
                else:
                    st.warning("No entities were found in the provided text.")
            
            end_time = time.time()
            elapsed_time = end_time - start_time
            st.text("")
            st.text("")
            st.info(f"Results processed in **{elapsed_time:.2f} seconds**.")

with tab2:
    

    # Load model for this tab
    model_qa = load_gliner_model("InfoFinder")

    user_text = st.text_area("Type or paste your text below, and then press Ctrl + Enter", height=250, key='my_text_area_infofinder')

    def clear_text_qa():
        st.session_state['my_text_area_infofinder'] = ""

    st.button("Clear text", on_click=clear_text_qa, key="clear_qa")

    st.subheader("Question-Answering", divider="green")
    question_input = st.text_input("Ask wh-questions. **Wh-questions begin with what, when, where, who, whom, which, whose, why and how. We use them to ask for specific information.**")

    if st.button("Add Question"):
        if question_input:
            if question_input not in st.session_state.user_labels:
                st.session_state.user_labels.append(question_input)
                st.success(f"Added question: {question_input}")
            else:
                st.warning("This question has already been added.")
        else:
            st.warning("Please enter a question.")
            
    st.markdown("---")
    st.subheader("Record of Questions", divider="green")

    if st.session_state.user_labels:
        for i, label in enumerate(st.session_state.user_labels):
            col_list, col_delete = st.columns([0.9, 0.1])
            with col_list:
                st.write(f"- {label}", key=f"label_{i}")
            with col_delete:
                if st.button("Delete", key=f"delete_{i}"):
                    st.session_state.user_labels.pop(i)
                    st.rerun()
    else:
        st.info("No questions defined yet. Use the input above to add one.")

    st.divider()

    if st.button("Extract Answers"):
        if not user_text.strip():
            st.warning("Please enter some text to analyze.")
        elif not st.session_state.user_labels:
            st.warning("Please define at least one question.")
        else:
            if comet_initialized:
                experiment = Experiment(api_key=COMET_API_KEY, workspace=COMET_WORKSPACE, project_name=COMET_PROJECT_NAME)
                experiment.log_parameter("input_text_length", len(user_text))
                experiment.log_parameter("defined_labels", st.session_state.user_labels)

            start_time = time.time()
            with st.spinner("Analyzing text...", show_time=True):
                try:
                    entities = model_qa.predict_entities(user_text, st.session_state.user_labels)
                    end_time = time.time()
                    elapsed_time = end_time - start_time
                    st.info(f"Processing took **{elapsed_time:.2f} seconds**.")

                    if entities:
                        df1 = pd.DataFrame(entities)
                        df2 = df1[['label', 'text', 'score']]
                        df = df2.rename(columns={'label': 'question', 'text': 'answer'})
                        
                        st.subheader("Extracted Answers", divider="green")
                        expander = st.expander("**Download**")
                        expander.write(""" 
                        To download the data, simply hover your cursor over the table. A download icon will appear on the right side.
                        """)
                        st.dataframe(df, use_container_width=True)

                        st.subheader("Tree map", divider="green")
                        all_labels = df['question'].unique()
                        label_color_map = {label: get_stable_color(label) for label in all_labels}
                        fig_treemap = px.treemap(df, path=[px.Constant("all"), 'question', 'answer'], values='score', color='question', color_discrete_map=label_color_map)
                        fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25), paper_bgcolor='#F3E5F5', plot_bgcolor='#F3E5F5')
                        st.plotly_chart(fig_treemap)

                       
                        if comet_initialized:
                            experiment.log_metric("processing_time_seconds", elapsed_time)
                            experiment.log_table("predicted_entities", df)
                            experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap")
                            experiment.end()
                    else:
                        st.info("No answers were found in the text with the defined questions.")
                        if comet_initialized:
                            experiment.end()
                except Exception as e:
                    st.error(f"An error occurred during processing: {e}")
                    st.write(f"Error details: {e}")
                    if comet_initialized:
                        experiment.log_text(f"Error: {e}")
                        experiment.end()