import streamlit as st
import joblib
import numpy as np

# Load the saved model and encoders
xgb_model = joblib.load('xgb_model.joblib')
label_encoder_X = joblib.load('label_encoder_X.joblib')
label_encoder_y = joblib.load('label_encoder_y.joblib')

# Define the Streamlit app
def main():
    st.title("Injury Type Prediction")

    # Get user input for features
    age = st.number_input("Age", min_value=0, max_value=120, step=1)
    sex = st.number_input("Sex", min_value=1, max_value=3)
    present_city = st.text_input("Present City")

    # Transform categorical inputs using encoders
    # Validate user input for present city
    if present_city not in label_encoder_X.classes_:
        st.warning("Unrecognized city. Please enter a valid city.")
        return

    # Transform inputs
    present_city_encoded = label_encoder_X.transform([present_city])[0]

    # Prepare input data
    input_data = np.array([[age, sex, present_city_encoded]])

    # Make predictions
    if st.button("Predict"):
        # Use the trained model to make predictions
        prediction = xgb_model.predict(input_data)

        # Decode the prediction
        predicted_injury_type = label_encoder_y.inverse_transform(prediction)
        if(predicted_injury_type[0]==2):
            st.write("Predicted Injury Type: Fatal")
        elif(predicted_injury_type[0]==3):
            st.write("Predicted Injury Type: Minor")
        elif(predicted_injury_type[0]==1):
            st.write("Predicted Injury Type: Grievous")
        elif(predicted_injury_type==4):
            st.write("Predicted Injury Type: Not Applicable")
        else:
            st.write("Predicted Injury Type: Abused")
        

if __name__ == "__main__":
    main()