import streamlit as st
import os
import requests
import pandas as pd
import json
import joblib  

# Function to create the payload for the model
def create_tf_serving_json(data):
    return {'inputs': {name: data[name].tolist() for name in data.keys()} if isinstance(data, dict) else data.tolist()}

# Function to send a request to the model endpoint
def score_model(dataset, url, token):
    headers = {'Authorization': f'Bearer {token}', 'Content-Type': 'application/json'}
    ds_dict = {'dataframe_split': dataset.to_dict(orient='split')} if isinstance(dataset, pd.DataFrame) else create_tf_serving_json(dataset)
    data_json = json.dumps(ds_dict, allow_nan=True)
    response = requests.request(method='POST', headers=headers, url=url, data=data_json)
    if response.status_code != 200:
        raise Exception(f'Request failed with status {response.status_code}, {response.text}')
    return response.json()

# Load local model function
def load_local_model():
    return joblib.load('Random Forest - Grid Search_model.pkl')

# Function to make predictions with local model
def predict_with_local_model(model, input_data):
    input_data = input_data[['satisfaction_level', 'last_evaluation', 'number_project', 
                             'average_montly_hours', 'time_spend_company', 'Work_accident', 
                             'promotion_last_5years', 'salaryVec_0', 'salaryVec_1']]
    prediction = model.predict(input_data)
    return prediction[0]

# Streamlit app UI
st.title('Employee Churn Prediction')

# Create a form
with st.form(key='churn_form'):
    satisfaction_level = st.slider('Satisfaction Level', 0.0, 1.0, 0.5)
    last_evaluation = st.slider('Last Evaluation', 0.0, 1.0, 0.5)
    number_project = st.slider('Number of Projects', 1, 10, 3)
    average_montly_hours = st.slider('Average Monthly Hours', 50, 350, 200)
    time_spend_company = st.slider('Time Spent in Company (years)', 1, 10, 3)
    work_accident = st.selectbox('Work Accident', [0, 1])
    promotion_last_5years = st.selectbox('Promotion in Last 5 Years', [0, 1])
    salary = st.selectbox('Salary Level', ['Low', 'Medium', 'High'])

    # Encode salary into one-hot vectors
    salaryVec_0, salaryVec_1 = 0.0, 0.0
    if salary == 'Low':
        salaryVec_0 = 1.0
    elif salary is 'Medium':
        salaryVec_1 = 1.0

    # Submit button
    submit_button = st.form_submit_button(label='Predict Churn')

# Handle form submission
if submit_button:
    # Create a DataFrame with the input data
    input_data = pd.DataFrame({
        'satisfaction_level': [satisfaction_level],
        'last_evaluation': [last_evaluation],
        'number_project': [number_project],
        'average_montly_hours': [average_montly_hours],
        'time_spend_company': [time_spend_company],
        'Work_accident': [work_accident],
        'promotion_last_5years': [promotion_last_5years],
        'salaryVec_0': [salaryVec_0],
        'salaryVec_1': [salaryVec_1]
    })

    # Check if URL and token are specified
    url = 'Add_Model_URL'
    token = 'Add_Databricks_Token'
    
    try:
        if url != 'Add_Model_URL' and token != 'Add_Databricks_Token':
            # Use the API if URL and token are specified
            prediction = score_model(input_data, url, token)
            churn_prediction = prediction['predictions'][0]
        else:
            # Use the local model if URL and token are not specified
            local_model = load_local_model()
            churn_prediction = predict_with_local_model(local_model, input_data)
        
        if churn_prediction == 1:
            st.write('The employee is likely to churn.')
        else:
            st.write('The employee is not likely to churn.')
    except Exception as e:
        st.error(f'Error: {e}')