import pandas as pd
import gradio as gr
import joblib
from gradio_modal import Modal


# Inputs for UI


def age_conversion(age):
    # Convert inputs to model inputs
    if 12 <= age <= 14:
        age_range = 1
    elif 15 <= age <= 17:
        age_range = 2
    elif 18 <= age <= 20:
        age_range = 3
    elif 21 <= age <= 24:
        age_range = 4
    elif 25 <= age <= 29:
        age_range = 5
    elif 30 <= age <= 34:
        age_range = 6
    elif 35 <= age <= 39:
        age_range = 7
    elif 40 <= age <= 44:
        age_range = 8
    elif 45 <= age <= 49:
        age_range = 9
    elif 50 <= age <= 54:
        age_range = 10
    elif 55 <= age <= 64:
        age_range = 11
    else:
        age_range = 12
    return age_range


def days_waiting_conversion(days_waiting):
    if days_waiting == 0:
        days_waiting_range = 0
    elif 1 <= days_waiting <= 7:
        days_waiting_range = 1
    elif 8 <= days_waiting <= 14:
        days_waiting_range = 2
    elif 15 <= days_waiting <= 30:
        days_waiting_range = 3
    else:
        days_waiting_range = 4

    return days_waiting_range


def first_use_conversion(first_use):

    if first_use <= 11:
        first_use_range = 1
    elif 12 <= first_use <= 14:
        first_use_range = 2
    elif 15 <= first_use <= 17:
        first_use_range = 3
    elif 18 <= first_use <= 20:
        first_use_range = 4
    elif 21 <= first_use <= 24:
        first_use_range = 5
    elif 25 <= first_use <= 29:
        first_use_range = 6
    else:
        first_use_range = 7
    return first_use_range


def load_model():
    # Load the model
    rf = joblib.load("random_forest_model.joblib")
    return rf


MODEL = load_model()


def call_the_model(
    age,
    arrests,
    substance,
    marital,
    education,
    primary_income,
    days_waiting,
    frequency,
    self_help,
    health_insurance,
    first_use,
    race,
    services,
    co_occuring,
):
    # create a dataframe
    data = {
        "AGE": [age_conversion(age)],
        "ARRESTS": [arrests],
        "SUB1": [substance + 1],
        "MARSTAT": [marital + 1],
        "EDUC": [education + 1],
        "PRIMINC": [primary_income + 1],
        "DAYWAIT": [days_waiting_conversion(days_waiting)],
        "FREQ1": [frequency + 1],
        "FREQ_ATND_SELF_HELP": [self_help + 1],
        "HLTHINS": [health_insurance + 1],
        "FRSTUSE1": [first_use_conversion(first_use)],
        "RACE": [race + 1],
        "SERVICES": [services + 1],
        "PSYPROB": [co_occuring],
    }
    input_df = pd.DataFrame(data)

    prob = MODEL.predict_proba(input_df)
    no, yes = prob[0]
    if no > yes:
        return f"The client is {round(no * 100)}% likely to relapse."
    else:
        return f"The client is {round(yes * 100)}% likely to not relapse."


def inference(age: gr.Number, race: gr.Dropdown):
    pass


with gr.Blocks(
    # theme=gr.themes.Soft(
    #     primary_hue=gr.themes.colors.zinc,
    #     neutral_hue=gr.themes.colors.slate,
    # )
) as demo:
    title = gr.Markdown("# NoLapse: Substance Use Relapse Prediction")
    with gr.Row():
        with gr.Column():
            age = gr.Number(
                label="Age",
                info="Age of client at admission",
                minimum=0,
                maximum=100,
                value=25,
            )
            race = gr.Dropdown(
                [
                    "Alaska Native (Aleut, Eskimo, Indian)",
                    "American Indian (other than Alaska Native)",
                    "Asian or Pacific Islander",
                    "Black or African American",
                    "White",
                    "Asian",
                    "Other single race",
                    "Two or more races",
                    "Native Hawaiian or Other Pacific Islander",
                ],
                label="Race",
                type="index",
                value=0,
            )
            education = gr.Dropdown(
                [
                    "Less than one school grade, no schooling, nursery school, or kindergarten to Grade 8",
                    "Grades 9 to 11",
                    "Grade 12 (or GED)",
                    "1-3 years of college, university, or vocational school",
                    "4 years of college, university, BA/BS, some postgraduate study, or more",
                ],
                label="Education",
                info="The highest school grade completed for adults or children not attending school, or current school grade for school-age children (3-17 years old) attending school",
                type="index",
                value=0,
            )
            marital = gr.Dropdown(
                ["Never married", "Now married", "Separated", "Divorced, widowed"],
                label="Marital Status",
                info="Client's marital status, compatible with U.S. Census categories",
                type="index",
                value=0,
            )
            primary_income = gr.Dropdown(
                [
                    "Wages/salary",
                    "Public assistance",
                    "Retirement/pension, disability",
                    "Other",
                    "None",
                ],
                label="Primary Source of Income/Support",
                info="Client's principal source of financial support (for children younger than 18 years old, the primary parental source of income/support)",
                type="index",
                value=0,
            )
            health_insurance = gr.Dropdown(
                [
                    "Private insurance, Blue Cross/Blue Shield, HMO",
                    "Medicaid",
                    "Medicare, other (e.g. TRICARE, CHAMPUS)",
                    "None",
                ],
                label="Health Insurance",
                info="Client's health insurance at admission",
                type="index",
                value=0,
            )
            frequency = gr.Radio(
                ["No use in the past month", "Some use", "Daily use"],
                label="Frequency of Primary Substance Use",
                type="index",
                value="No use in the past month",
            )

        with gr.Column():
            primary_substance = gr.Dropdown(
                [
                    "None",
                    "Alcohol",
                    "Cocaine/crack",
                    "Marijuana/hashish: Includes THC and any other cannabis sativa preparations",
                    "Heroin",
                    "Non-prescription methadone",
                    "Other opiates and synthetics: Includes buprenorphine, butorphanol, codeine, hydrocodone, hydromorphone, meperidine,morphine, opium, oxycodone, pentazocine, propoxyphene, tramadol, and other narcotic analgesics, opiates, or synthetics",
                    "PCP: Phencyclidine",
                    "Hallucinogens: Includes LSD, DMT, mescaline, peyote, psilocybin, STP, and other hallucinogens",
                    "Methamphetamine/speed",
                    "Other amphetamines: Includes amphetamines, MDMA, ‘bath salts’, phenmetrazine, and other amines and related drugs",
                    "Other stimulants: Includes methylphenidate and any other stimulants",
                    "Benzodiazepines: Includes alprazolam, chlordiazepoxide, clonazepam, clorazepate, diazepam, flunitrazepam,flurazepam, halazepam, lorazepam, oxazepam, prazepam, temazepam, triazolam, and other unspecified benzodiazepines",
                    "Other tranquilizers: Includes meprobamate, and other non-benzodiazepine tranquilizers",
                    "Barbiturates: Includes amobarbital, pentobarbital, phenobarbital, secobarbital, etc.",
                    "Other sedatives or hypnotics: Includes chloral hydrate, ethchlorvynol, glutethimide, methaqualone, and othernon-barbiturate sedatives and hypnotics",
                    "Inhalants: Includes aerosols; chloroform, ether, nitrous oxide and other anesthetics; gasoline; glue; nitrites; paint thinnerand other solvents; and other inappropriately inhaled products",
                    "Over-the-counter medications: Includes aspirin, dextromethorphan and other cough syrups, diphenhydramine and otheranti-histamines, ephedrine, sleep aids, and any other legally obtained, non-prescription medication",
                    "Other drugs: Includes diphenylhydantoin/phenytoin, GHB/GBL, ketamine, synthetic cannabinoid 'Spice', carisoprodol(Soma), and other drugs",
                ],
                label="Primary Substance Use",
                info="Client's primary substance use at admission",
                type="index",
                value=0,
            )

            first_use = gr.Number(
                label="Age at First Use of Primary Substance",
                info="For alcohol use, this is the age of first intoxication",
                minimum=0,
                maximum=100,
                value=25,
            )

            days_waiting = gr.Number(
                label="Days Waiting to Enter Substance Use Treatment",
                info="Number of days from the first contact or request for substance use treatment service until the client was admitted and the first clinical substance use treatment service was provided",
                minimum=0,
            )
            arrests = gr.Radio(
                ["None", "Once", "Two or more times"],
                label="Arrests",
                info="Number of arrests in the 30 days prior to admission",
                type="index",
                value="None",
            )
            attendance = gr.Dropdown(
                [
                    "No attendance",
                    "1-3 times in the past month",
                    "4-7 times in the past month",
                    "8-30 times in the past month",
                    "Some attendance, frequency is unknown",
                ],
                label="Attendance at Substance Use Self-help Groups in Past 30 Days",
                info="Frequency of attendance at a substance use self-help group in the 30 days prior to the reference date (the date of admission). Includes Alcoholics Anonymous (AA), Narcotics Anonymous (NA), and other self-help/mutual support groups focused on recovery from substance use and dependence",
                type="index",
                value=0,
            )
            services = gr.Dropdown(
                [
                    "Detox, 24-hour, hospital inpatient",
                    "Detox, 24-hour, free-standing residential",
                    "Rehab/residential, hospital (non-detox)",
                    "Rehab/residential, short term (30 days or fewer)",
                    "Rehab/residential, long term (more than 30 days)",
                    "Ambulatory, intensive outpatient",
                    "Ambulatory, non-intensive outpatient",
                    "Ambulatory, detoxification",
                ],
                label="Type of Treatment Service/Setting",
                info="Type of treatment service or treatment setting in which the client is placed at the time of admission or transfer",
                type="index",
                value=0,
            )

            co_occuring = gr.Checkbox(
                value=False, label="Co-occurring Mental and Substance Use Disorders"
            )

    submit_btn = gr.Button("SUBMIT", variant="primary", size="lg")

    # result = gr.Label(label="Result")
    with Modal(visible=False) as modal:
        modal_content = gr.Label(show_label=False)
        disclaimer = gr.Markdown(
            "###### The accuracy of this model is 74%. Take results with caution."
        )

    submit_btn.click(
        call_the_model,
        [
            age,
            arrests,
            primary_substance,
            marital,
            education,
            primary_income,
            days_waiting,
            frequency,
            attendance,
            health_insurance,
            first_use,
            race,
            services,
            co_occuring,
        ],
        [modal_content],
    ).then(lambda: Modal(visible=True), None, modal)


if __name__ == "__main__":
    demo.launch()