import gradio as gr
import pickle

# Model URL for each
lr_model_url = './logistic regression/logistic_regression_model.pkl'

def bool_value(val):
    if val:
        return 1
    else:
        return 0


# Make Prediction Model but would also like to add Gender and Race
def make_prediction(personal_injury, property_damage, fatal, commercial_vehicle, alcohol, rsam, disobedience, invalid_docu, phone, speeding, 
                    negligent, vss, num_offences):

    # load model
    with open(lr_model_url, 'rb') as file:  
        lr_model = pickle.load(file)

    x_input_feature = [[bool_value(personal_injury), bool_value(property_damage), bool_value(fatal), bool_value(commercial_vehicle), bool_value(alcohol),
                        bool_value(rsam), bool_value(disobedience), bool_value(invalid_docu), bool_value(phone), bool_value(speeding), bool_value(negligent),
                        bool_value(vss), num_offences]]
    prd = lr_model.predict(x_input_feature)

    if prd == 0:
        return 'SERO'
    elif prd == 1:
        return 'Warning'
    else:
        return 'Citation'


iface = gr.Interface(fn=make_prediction,
                     inputs=[gr.components.Checkbox(label='Did the violation involve any personal injury?'),
                             gr.components.Checkbox(label='Did the violation involve any property damage?'),
                             gr.components.Checkbox(label='Did the violation involve any fatalities?'),
                             gr.components.Checkbox(label='Is the vehicle committing the traffic violation a commercial vehicle?'),
                             gr.components.Checkbox(label='Was the driver under the influence of alcohol?'),
                             gr.components.Checkbox(label='Did the driver fail to obey signs and markings (such as traffic control device instructions, stop lights, red signal and stop sign lines)?'),
                             gr.components.Checkbox(label='Was the driver disobedient? (such as failing to display documentation upon request)?'),
                             gr.components.Checkbox(label='Was the driver driving with Invalid Documentation (such as suspended registration, suspended license, expired registration plates and validation tabs or expired license plate)?'),
                             gr.components.Checkbox(label='Was the driver using a mobile phone while driving?'),
                             gr.components.Checkbox(label='Was the driver caught speeding?'),
                             gr.components.Checkbox(label='Was the driver caught driving with negligence (example switching lanes in an unsafe manner)?'),
                             gr.components.Checkbox(label='Was the vehicle up to standards (lights properly switched, registration plates attached etc.)?'),
                             gr.components.Slider(minimum=1, step=1, label='Number of offences committed')],
                     outputs=["text"])

iface.launch(debug=True)