## Importing the libraries
import joblib
import gradio as gr

# Loading the trained model
random_forest_model = joblib.load('tuned_random_forest_model.joblib')
xgb_model = joblib.load('xgb_model_release.joblib')
logistic_regression_model = joblib.load('lr_lbfgs_model.joblib')

def random_forest_predict(new_id, judge_id, prior_misdemeanor, race, county, med_house_income, highest_severity, probation):
    race_to_num = {'African American': 0, 'American Indian or Alaskan Native': 1, 'Asian or Pacific Islander': 2, 'Caucasian': 3, 'Hispanic': 4}
    race_num = race_to_num[race]
    processed_input = [new_id, judge_id, prior_misdemeanor, race_num, county, med_house_income, highest_severity, probation]

    # Predicting
    prediction = random_forest_model.predict([processed_input])
    return "The prisoner is likely to return to jail." if prediction[0] == 1 else "The prisoner is unlikely to return to jail."

# Input components
random_forest_input_components = [
    gr.Number(label="Prisoner ID", value=0),
    gr.Number(label="Judge ID", value=0),
    gr.Number(label="Prior Misdemeanor", value=0),
    gr.Dropdown(label="Race", choices={'African American': 0, 'American Indian or Alaskan Native': 1, 'Asian or Pacific Islander': 2, 'Caucasian': 3, 'Hispanic': 4}, value='Caucasian'),
    gr.Number(label="County", value=0),
    gr.Number(label="Median House Income", value=0),
    gr.Number(label="Highest Severity", value=0),
    gr.Number(label="Probation", value=0)
]

random_forest_interface = gr.Interface(
    fn=random_forest_predict,
    inputs=random_forest_input_components,
    outputs="text",
    title="Random Forest Prediction",
    description="Predicts the likelihood of a prisoner returning to jail using a Random Forest model."
)

def xgb_predict(new_id, judge_id, sex, race, prior_felony, highest_severity, age_judge, probation):
    race_to_num = {'African American': 0, 'American Indian or Alaskan Native': 1, 'Asian or Pacific Islander': 2, 'Caucasian': 3, 'Hispanic': 4}
    race_num = race_to_num[race]
    sex_to_num = {'F': 0, 'M': 1}
    sex_num = sex_to_num[sex]
    processed_input = [new_id, judge_id, sex_num, race_num, prior_felony, highest_severity, age_judge, probation]

    # Predicting
    prediction = xgb_model.predict([processed_input])
    return "The defendant is unlikely to go to jail." if prediction[0] == 1 else "The defendant is likely to go to jail."

# Input components
xgb_input_components = [
    gr.Number(label="Prisoner ID", value=0),
    gr.Number(label="Judge ID", value=0),
    gr.Dropdown(label="Sex", choices={'F': 0, 'M': 1}, value='M'),
    gr.Dropdown(label="Race", choices={'African American': 0, 'American Indian or Alaskan Native': 1, 'Asian or Pacific Islander': 2, 'Caucasian': 3, 'Hispanic': 4}, value='Caucasian'),
    gr.Number(label="Prior Felony", value=0),
    gr.Number(label="Highest Severity", value=0),
    gr.Number(label="Age of Judge", value=0),
    gr.Number(label="Probation", value=0)
]

xgb_interface = gr.Interface(
    fn=xgb_predict,
    inputs=xgb_input_components,
    outputs="text",
    title="XGBoost Prediction",
    description="Predicts the likelihood of a prisoner returning to jail using an XGBoost model."
)

def logistic_regression_predict(new_id, sex, race, offence_category, age_judge, prior_felony, probation, year, case_type):
    race_to_num = {'African American': 0, 'American Indian or Alaskan Native': 1, 'Asian or Pacific Islander': 2, 'Caucasian': 3, 'Hispanic': 4}
    race_num = race_to_num[race]
    case_type_to_num = {'Criminal Traffic': 0, 'Felony': 1, 'Misdemeanor': 2}
    case_type_num = case_type_to_num[case_type]
    offence_category_to_num = {'1st Deg. Sex. Assault of Child': 0, '2nd Deg. Sex. Assault of Child': 1, 'Armed Robbery': 2, 'Arson': 3, 'BAC': 4, 'Bail Jumping': 5, 'Battery': 6, 'Burglary': 7, 'Child Abuse': 8, 'Contempt of Court': 9, 'Crimes Against Children': 10, 'Criminal Damage': 11, 'Criminal Trespass': 12, 'Disorderly Conduct': 13, 'Drug Manufacture/Deliver': 14, 'Drug Paraphernalia': 15, 'Drug Possession': 16, 'Entering Locked Vehicle': 17, 'Escape': 18, 'Extradition': 19, 'First Degree Intentional Homicide': 20, 'First Degree Reckless Homicide': 21, 'Fleeing/Eluding': 22, 'Forgery': 23, 'Fourth Degree Sexual Assau': 24, 'Gambling': 25, 'Hit and Run': 26, 'Intimidate Witness/Victim': 27, 'Kidnap/Hostage/False Imprisonment': 28, 'Local or Unidentified Forfeiture': 29, 'Non-Traffic Forfeiture': 30, 'OAR/OAS': 31, 'Operate Vehicle Without Consent': 32, 'Operate Vehicle w/out Consent': 33, 'Operate Without License': 34, 'Operating While Intoxicated': 35, 'Operating while intoxicated': 36, 'Other Bodily Security': 37, 'Other Crimes Against Children': 38, 'Other Drug Offenses': 39, 'Other Felony': 40, 'Other Fraud': 41, 'Other Homicide': 42, 'Other Misdemeanor': 43, 'Other Public Safety Crimes': 44, 'Perjury': 45, 'Public Assistance Fraud': 46, 'Receiving Stolen Property': 47, 'Reckless Driving': 48, 'Resisting Officer': 49, 'Retail Theft (Shoplifting)': 50, 'Sex Crimes': 51, 'Sexual Assault': 52, 'Stalking': 53, 'Substantial/Aggravated Battery': 54, 'Theft': 55, 'Unarmed Robbery': 56, 'Unidentified Felony': 57, 'Unidentified Felony Traffic': 58, 'Unidentified Misdemeanor': 59, 'Unidentified Misdemeanor Traffic': 60, 'Unidentified Traffic Forfeiture': 61, 'Violate Occupational': 62, 'Violation of TRO': 63, 'Weapons/Explosives': 64, 'Worthless Checks': 65}
    offence_category_num = offence_category_to_num[offence_category]
    sex_to_num = {'F': 0, 'M': 1}
    sex_num = sex_to_num[sex]
    
    processed_input = [new_id, sex_num, race_num, offence_category_num, age_judge, prior_felony, probation, year, case_type_num]

    # Predicting
    prediction = logistic_regression_model.predict([processed_input])
    return "The defendant is likely to be detained (not released)." if prediction[0] == 1 else "The defendant is likely to be restrained."

# Input components
logistic_regression_input_components = [
    gr.Number(label="Prisoner ID", value=0),
    gr.Dropdown(label="Sex", choices={'F': 0, 'M': 1}, value='M'),
    gr.Dropdown(label="Race", choices={'African American': 0, 'American Indian or Alaskan Native': 1, 'Asian or Pacific Islander': 2, 'Caucasian': 3, 'Hispanic': 4}, value='Caucasian'),
    gr.Dropdown(label="Offence Category", choices={'1st Deg. Sex. Assault of Child': 0, '2nd Deg. Sex. Assault of Child': 1, 'Armed Robbery': 2, 'Arson': 3, 'BAC': 4, 'Bail Jumping': 5, 'Battery': 6, 'Burglary': 7, 'Child Abuse': 8, 'Contempt of Court': 9, 'Crimes Against Children': 10, 'Criminal Damage': 11, 'Criminal Trespass': 12, 'Disorderly Conduct': 13, 'Drug Manufacture/Deliver': 14, 'Drug Paraphernalia': 15, 'Drug Possession': 16, 'Entering Locked Vehicle': 17, 'Escape': 18, 'Extradition': 19, 'First Degree Intentional Homicide': 20, 'First Degree Reckless Homicide': 21, 'Fleeing/Eluding': 22, 'Forgery': 23, 'Fourth Degree Sexual Assau': 24, 'Gambling': 25, 'Hit and Run': 26, 'Intimidate Witness/Victim': 27, 'Kidnap/Hostage/False Imprisonment': 28, 'Local or Unidentified Forfeiture': 29, 'Non-Traffic Forfeiture': 30, 'OAR/OAS': 31, 'Operate Vehicle Without Consent': 32, 'Operate Vehicle w/out Consent': 33, 'Operate Without License': 34, 'Operating While Intoxicated': 35, 'Operating while intoxicated': 36, 'Other Bodily Security': 37, 'Other Crimes Against Children': 38, 'Other Drug Offenses': 39, 'Other Felony': 40, 'Other Fraud': 41, 'Other Homicide': 42, 'Other Misdemeanor': 43, 'Other Public Safety Crimes': 44, 'Perjury': 45, 'Public Assistance Fraud': 46, 'Receiving Stolen Property': 47, 'Reckless Driving': 48, 'Resisting Officer': 49, 'Retail Theft (Shoplifting)': 50, 'Sex Crimes': 51, 'Sexual Assault': 52, 'Stalking': 53, 'Substantial/Aggravated Battery': 54, 'Theft': 55, 'Unarmed Robbery': 56, 'Unidentified Felony': 57, 'Unidentified Felony Traffic': 58, 'Unidentified Misdemeanor': 59, 'Unidentified Misdemeanor Traffic': 60, 'Unidentified Traffic Forfeiture': 61, 'Violate Occupational': 62, 'Violation of TRO': 63, 'Weapons/Explosives': 64, 'Worthless Checks': 65}, value='Worthless Checks'),
    gr.Number(label="Age of Judge", value=0),
    gr.Number(label="Prior Felony", value=0),
    gr.Number(label="Probation", value=0),
    gr.Number(label="Year", value=0),
    gr.Dropdown(label="Case Type", choices={'Criminal Traffic': 0, 'Felony': 1, 'Misdemeanor': 2}, value='Felony')
]

logistic_regression_interface = gr.Interface(
    fn=logistic_regression_predict,
    inputs=logistic_regression_input_components,
    outputs="text",
    title="Logistic Regression Prediction",
    description="Predicts the likelihood of a defendant being detained or released using a Logistic Regression model."
)

tabbed_interface = gr.TabbedInterface(
    [random_forest_interface, xgb_interface, logistic_regression_interface],
    ["Random Forest", "XGBoost", "Logistic Regression"]
)
tabbed_interface.launch(share=True)