import streamlit as st
from huggingface_hub import hf_hub_download
from model import JobFakeModel
import torch
import torch.nn.functional as F
from data_processing import get_tokens
import numpy as np
import time

def get_model_list():
    model_name = "sebastiansarasti/fakeJobs"  
    filename = "best_model.pth"  
    file_path = hf_hub_download(repo_id=model_name, filename=filename)
    return file_path

def load_model(path):
    model = JobFakeModel(base_model="distilbert", freeze_base=True)
    model.load_state_dict(torch.load(path, map_location=torch.device('cpu')))
    return model

st.title('Fake Jobs Streamlit App')
st.write('This is a fake jobs streamlit app')

# download the model
path = get_model_list()
model = load_model(path)

with st.sidebar:
    st.subheader('About the App')
    st.markdown('Data used for the training come from the following source: https://www.kaggle.com/datasets/shivamb/real-or-fake-fake-jobposting-prediction')
    st.empty()
    st.subheader('Author')
    st.markdown('Sebastián Sarasti Zambonino')
    st.markdown('Data Scientist - Machine Learning Engineer')
    st.markdown('https://www.linkedin.com/in/sebastiansarasti/')
    st.markdown('https://github.com/sebassaras02')

# Create columns for the inputs
col1, col2, col3 = st.columns(3)

with col1:
    # Create an input text box for the description
    description = st.text_area('Description', 'Enter the job description here')

with col2:
    # Create an input text box for the requirements
    requirements = st.text_area('Requirements', 'Enter the job requirements here')

with col3:
    # Create an input text box for the benefits
    benefits = st.text_area('Benefits', 'Enter the job benefits here')

# if benefits is none, set it to an empty nothing
if benefits is None:
    benefits = 'Nothing'
elif requirements is None:
    requirements = 'Nothing'
elif description is None:
    raise ValueError('Description cannot be empty')

# Create a button to submit the job
with st.spinner('Wait for it...'):
    time.sleep(2)
    if st.button('Submit Job'):
        tokens_des = get_tokens(description)
        tokens_req = get_tokens(requirements)
        tokens_ben = get_tokens(benefits)
        model.eval()
        with torch.no_grad():
            output = model(tokens_des, tokens_req, tokens_ben)
        # calculate the probability
        output = F.sigmoid(output).item()*100
        # Create a box to show the result
        st.metric('Fake Job Probability', f"{np.round(output, 5)}%")
        if output > 75:
            st.error('⚠️ High probability of being a fake job!')
        elif output > 50:
            st.warning('⚠️ Medium probability of being a fake job')
        else:
            st.success('✅ Low probability of being a fake job')