app.py

import pickle

import gradio as gr
import numpy as np
import pandas as pd
import plotly.express as px

# Load the training CSV once (outside the functions so it is read only once).
df = pd.read_csv("X_train_Y_Train_merged_train.csv")

###############################################################################
# 1) Model Predictor class
###############################################################################
class ModelPredictor:
    def __init__(self, model_path, model_filenames):
        self.model_path = model_path
        self.model_filenames = model_filenames
        self.models = self.load_models()

        # For each model name, define the mapping from 0->..., 1->...
        # If you have more labels, expand this dictionary accordingly.
        self.prediction_map = {
            "YOWRCONC": ["Did not have difficulty concentrating", "Had difficulty concentrating"],
            "YOSEEDOC": ["Did not feel the need to see a doctor", "Felt the need to see a doctor"],
            "YOWRHRS": ["Did not have trouble sleeping", "Had trouble sleeping"],
            "YO_MDEA5": ["Others did not notice restlessness/lethargy", "Others noticed restlessness/lethargy"],
            "YOWRCHR": ["Did not feel so sad that nothing could cheer up", "Felt so sad that nothing could cheer up"],
            "YOWRLSIN": ["Did not feel bored and lose interest in all enjoyable things", 
                         "Felt bored and lost interest in all enjoyable things"],
            "YODPPROB": ["Did not have other problems for 2+ weeks", "Had other problems for 2+ weeks"],
            "YOWRPROB": ["Did not have the worst time ever feeling", "Had the worst time ever feeling"],
            "YODPR2WK": ["Did not have periods where feelings lasted 2+ weeks", 
                         "Had periods where feelings lasted 2+ weeks"],
            "YOWRDEPR": ["Did not feel sad/depressed mostly everyday", "Felt sad/depressed mostly everyday"],
            "YODPDISC": ["Overall mood duration was not sad/depressed", 
                         "Overall mood duration was sad/depressed (discrepancy)"],
            "YOLOSEV": ["Did not lose interest in enjoyable things and activities",
                        "Lost interest in enjoyable things and activities"],
            "YOWRDCSN": ["Was able to make decisions", "Was unable to make decisions"],
            "YODSMMDE": ["Never had depression symptoms lasting 2 weeks or longer",
                         "Had depression symptoms lasting 2 weeks or longer"],
            "YO_MDEA3": ["Did not experience changes in appetite or weight",
                         "Experienced changes in appetite or weight"],
            "YODPLSIN": ["Never lost interest and felt bored", "Lost interest and felt bored"],
            "YOWRELES": ["Did not eat less than usual", "Ate less than usual"],
            "YODSCEV": ["Had fewer severe symptoms of depression", "Had more severe symptoms of depression"],
            "YOPB2WK": ["Did not experience uneasy feelings lasting every day for 2+ weeks or longer",
                        "Experienced uneasy feelings lasting every day for 2+ weeks or longer"],
            "YO_MDEA2": ["Did not have issues with physical and mental well-being every day for 2 weeks or longer",
                         "Had issues with physical and mental well-being every day for 2 weeks or longer"]
        }

    def load_models(self):
        models = []
        for filename in self.model_filenames:
            filepath = self.model_path + filename
            with open(filepath, 'rb') as file:
                model = pickle.load(file)
            models.append(model)
        return models

    def make_predictions(self, user_input):
        """
        Returns a list of numpy arrays, each array is [0] or [1].
        The i-th array corresponds to the i-th model in self.models.
        """
        predictions = []
        for model in self.models:
            pred = model.predict(user_input)
            pred = np.array(pred).flatten()
            predictions.append(pred)
        return predictions

    def get_majority_vote(self, predictions):
        """
        Flatten all predictions from all models, combine them into a single array,
        then find the majority class (0 or 1) across all of them.
        """
        combined_predictions = np.concatenate(predictions)
        majority_vote = np.bincount(combined_predictions).argmax()
        return majority_vote

    # Based on Equal Interval and Percentage-Based Method 
    # Severe: 13 to 16 votes (upper 25%) 
    # Moderate: 9 to 12 votes (upper-middle 25%) 
    # Low: 5 to 8 votes (lower-middle 25%) 
    # Very Low: 0 to 4 votes (lower 25%)
    def evaluate_severity(self, majority_vote_count):
        if majority_vote_count >= 13:
            return "Mental health severity: Severe"
        elif majority_vote_count >= 9:
            return "Mental health severity: Moderate"
        elif majority_vote_count >= 5:
            return "Mental health severity: Low"
        else:
            return "Mental health severity: Very Low"

###############################################################################
# 2) Model Filenames & Predictor
###############################################################################
model_filenames = [
    "YOWRCONC.pkl", "YOSEEDOC.pkl", "YO_MDEA5.pkl", "YOWRLSIN.pkl",
    "YODPPROB.pkl", "YOWRPROB.pkl", "YODPR2WK.pkl", "YOWRDEPR.pkl",
    "YODPDISC.pkl", "YOLOSEV.pkl", "YOWRDCSN.pkl", "YODSMMDE.pkl",
    "YO_MDEA3.pkl", "YODPLSIN.pkl", "YOWRELES.pkl", "YOPB2WK.pkl"
]
model_path = "models/"
predictor = ModelPredictor(model_path, model_filenames)

###############################################################################
# 3) Validate Inputs
###############################################################################
def validate_inputs(*args):
    for arg in args:
        if arg == '' or arg is None:  # Assuming empty string or None as unselected
            return False
    return True

###############################################################################
# 4) Reverse Lookup (numeric -> user-friendly text) for input columns
###############################################################################
# We'll define the forward mapping here. The reverse mapping is constructed below.
input_mapping = {
    'YNURSMDE': {"Yes": 1, "No": 0},
    'YMDEYR': {"Yes": 1, "No": 2},
    'YSOCMDE': {"Yes": 1, "No": 0},
    'YMDESUD5ANYO': {"SUD only, no MDE": 1, "MDE only, no SUD": 2, "SUD and MDE": 3, "Neither SUD or MDE": 4},
    'YMSUD5YANY': {"Yes": 1, "No": 0},
    'YUSUITHK': {"Yes": 1, "No": 2, "I'm not sure": 3, "I don't want to answer": 4},
    'YMDETXRX': {"Yes": 1, "No": 0},
    'YUSUITHKYR': {"Yes": 1, "No": 2, "I'm not sure": 3, "I don't want to answer": 4},
    'YMDERSUD5ANY': {"Yes": 1, "No": 0},
    'YUSUIPLNYR': {"Yes": 1, "No": 2, "I'm not sure": 3, "I don't want to answer": 4},
    'YCOUNMDE': {"Yes": 1, "No": 0},
    'YPSY1MDE': {"Yes": 1, "No": 0},
    'YHLTMDE': {"Yes": 1, "No": 0},
    'YDOCMDE': {"Yes": 1, "No": 0},
    'YPSY2MDE': {"Yes": 1, "No": 0},
    'YMDEHARX': {"Yes": 1, "No": 0},
    'LVLDIFMEM2': {"No Difficulty": 1, "Some difficulty": 2, "A lot of difficulty or cannot do at all": 3},
    'MDEIMPY': {"Yes": 1, "No": 2},
    'YMDEHPO': {"Yes": 1, "No": 0},
    'YMIMS5YANY': {"Yes": 1, "No": 0},
    'YMDEIMAD5YR': {"Yes": 1, "No": 0},
    'YMIUD5YANY': {"Yes": 1, "No": 0},
    'YMDEHPRX': {"Yes": 1, "No": 0},
    'YMIMI5YANY': {"Yes": 1, "No": 0},
    'YUSUIPLN': {"Yes": 1, "No": 2, "I'm not sure": 3, "I don't want to answer": 4},
    'YTXMDEYR': {"Yes": 1, "No": 0},
    'YMDEAUD5YR': {"Yes": 1, "No": 0},
    'YRXMDEYR': {"Yes": 1, "No": 0},
    'YMDELT': {"Yes": 1, "No": 2}
}

# Build reverse mapping: { "YNURSMDE": {1: "Yes", 0: "No"}, ... } etc.
reverse_mapping = {}
for col, mapping_dict in input_mapping.items():
    rev = {v: k for k, v in mapping_dict.items()}  # invert dict
    reverse_mapping[col] = rev

###############################################################################
# 5) Main Predict Function
###############################################################################
def predict(
    YMDEYR, YMDERSUD5ANY, YMDEIMAD5YR, YMIMS5YANY, YMDELT, YMDEHARX, 
    YMDEHPRX, YMDETXRX, YMDEHPO, YMDEAUD5YR, YMIMI5YANY, YMIUD5YANY, 
    YMDESUD5ANYO, YNURSMDE, YSOCMDE, YCOUNMDE, YPSY1MDE, YPSY2MDE, 
    YHLTMDE, YDOCMDE, YTXMDEYR, YUSUITHKYR, YUSUIPLNYR, YUSUITHK, 
    YUSUIPLN, MDEIMPY, LVLDIFMEM2, YMSUD5YANY, YRXMDEYR
):
    """
    Core prediction function that:
      1) Predicts with each model
      2) Aggregates results
      3) Produces an overall 'severity'
      4) Returns detailed per-model predictions
      5) Creates a distribution plot for ALL input features vs. a chosen label
      6) Nearest neighbor logic (with disclaimers), mapping numeric -> user text
    """

    # 1) Prepare user_input dataframe
    user_input_data = {
        'YNURSMDE': [int(YNURSMDE)],
        'YMDEYR': [int(YMDEYR)],
        'YSOCMDE': [int(YSOCMDE)],
        'YMDESUD5ANYO': [int(YMDESUD5ANYO)],
        'YMSUD5YANY': [int(YMSUD5YANY)],
        'YUSUITHK': [int(YUSUITHK)],
        'YMDETXRX': [int(YMDETXRX)],
        'YUSUITHKYR': [int(YUSUITHKYR)],
        'YMDERSUD5ANY': [int(YMDERSUD5ANY)],
        'YUSUIPLNYR': [int(YUSUIPLNYR)],
        'YCOUNMDE': [int(YCOUNMDE)],
        'YPSY1MDE': [int(YPSY1MDE)],
        'YHLTMDE': [int(YHLTMDE)],
        'YDOCMDE': [int(YDOCMDE)],
        'YPSY2MDE': [int(YPSY2MDE)],
        'YMDEHARX': [int(YMDEHARX)],
        'LVLDIFMEM2': [int(LVLDIFMEM2)],
        'MDEIMPY': [int(MDEIMPY)],
        'YMDEHPO': [int(YMDEHPO)],
        'YMIMS5YANY': [int(YMIMS5YANY)],
        'YMDEIMAD5YR': [int(YMDEIMAD5YR)],
        'YMIUD5YANY': [int(YMIUD5YANY)],
        'YMDEHPRX': [int(YMDEHPRX)],
        'YMIMI5YANY': [int(YMIMI5YANY)],
        'YUSUIPLN': [int(YUSUIPLN)],
        'YTXMDEYR': [int(YTXMDEYR)],
        'YMDEAUD5YR': [int(YMDEAUD5YR)],
        'YRXMDEYR': [int(YRXMDEYR)],
        'YMDELT': [int(YMDELT)]
    }
    user_input = pd.DataFrame(user_input_data)

    # 2) Make predictions
    predictions = predictor.make_predictions(user_input)

    # 3) Calculate majority vote (0 or 1) across all models
    majority_vote = predictor.get_majority_vote(predictions)

    # 4) Count how many 1's in all predictions combined
    majority_vote_count = sum([1 for pred in np.concatenate(predictions) if pred == 1])

    # 5) Evaluate severity
    severity = predictor.evaluate_severity(majority_vote_count)

    # 6) Prepare per-model predictions
    #    We'll group them just like before
    results = {
        "Concentration_and_Decision_Making": [],
        "Sleep_and_Energy_Levels": [],
        "Mood_and_Emotional_State": [],
        "Appetite_and_Weight_Changes": [],
        "Duration_and_Severity_of_Depression_Symptoms": []
    }
    
    prediction_groups = {
        "Concentration_and_Decision_Making": ["YOWRCONC", "YOWRDCSN"],
        "Sleep_and_Energy_Levels": ["YOWRHRS", "YO_MDEA5", "YOWRELES", "YO_MDEA2"],
        "Mood_and_Emotional_State": ["YOWRCHR", "YOWRLSIN", "YOWRDEPR", "YODPDISC", 
                                     "YOLOSEV", "YODPLSIN", "YODSCEV"],
        "Appetite_and_Weight_Changes": ["YO_MDEA3", "YOWRELES"],
        "Duration_and_Severity_of_Depression_Symptoms": ["YODPPROB", "YOWRPROB", 
                                                         "YODPR2WK", "YODSMMDE", 
                                                         "YOPB2WK"]
    }

    # We'll keep a record of which model => which predicted label
    for i, pred in enumerate(predictions):
        model_name = predictor.model_filenames[i].split('.')[0]
        pred_value = pred[0]
        # Map the prediction value to a human-readable string
        if model_name in predictor.prediction_map and pred_value in [0, 1]:
            result_text = f"Model {model_name}: {predictor.prediction_map[model_name][pred_value]}"
        else:
            result_text = f"Model {model_name}: Unknown or out-of-range"

        # Append to the appropriate group
        found_group = False
        for group_name, group_models in prediction_groups.items():
            if model_name in group_models:
                results[group_name].append(result_text)
                found_group = True
                break
        if not found_group:
            # If no group matches, skip or store in "Other"
            pass

    # 7) Nicely format the results
    formatted_results = []
    for group, preds in results.items():
        if preds:
            formatted_results.append(f"Group {group.replace('_', ' ')}:")
            formatted_results.append("\n".join(preds))
            formatted_results.append("\n")
    formatted_results = "\n".join(formatted_results).strip()
    if len(formatted_results) == 0:
        formatted_results = "No predictions made. Please check your inputs."

    # 8) Additional disclaimers if there's a large fraction of unknown
    num_unknown = sum(1 for group, preds in results.items() if any("Unknown or out-of-range" in p for p in preds))
    if num_unknown > len(predictor.model_filenames) / 2:
        severity += " (Unknown prediction count is high. Please consult with a human.)"

    ############################################################################
    # A) Total Patient Count
    ############################################################################
    total_patients = len(df)
    total_patient_count_markdown = (
        "### Total Patient Count\n"
        f"There are **{total_patients}** total patients in the dataset.\n\n"
        "This number helps you understand the size of the dataset used."
    )

    ############################################################################
    # B) Distribution Plot: All Input Features vs. a single predicted label
    ############################################################################
    # For demonstration, let's pick "YOWRCONC" if it exists in df:
    # We'll melt the dataset so that each input feature is in a "FeatureName" column,
    # and each distinct category is in "FeatureValue". We'll group by those + label to get counts.
    chosen_label = "YOWRCONC"
    if chosen_label in df.columns:
        # 1) Narrow down to the columns of interest
        #    We'll only use the input features that exist in df
        input_cols_in_df = [c for c in user_input_data.keys() if c in df.columns]
        # 2) We'll create a "melted" version of these input features
        #    i.e., row per (patient_id, FeatureName, FeatureValue)
        sub_df = df[input_cols_in_df + [chosen_label]].copy()
        # Melt them
        melted = sub_df.melt(
            id_vars=[chosen_label], 
            var_name="FeatureName", 
            value_name="FeatureValue"
        )
        # 3) Group by (FeatureName, FeatureValue, chosen_label) to get size
        dist_data = melted.groupby(["FeatureName", "FeatureValue", chosen_label]).size().reset_index(name="count")
        # 4) We'll try to map FeatureValue from numeric -> user-friendly text if possible
        #    We'll do it only if FeatureName is in reverse_mapping.
        def map_value(row):
            fn = row["FeatureName"]
            fv = row["FeatureValue"]
            if fn in reverse_mapping:
                if fv in reverse_mapping[fn]:
                    return reverse_mapping[fn][fv]  # e.g. 1->"Yes"
            return fv  # fallback
        dist_data["FeatureValueText"] = dist_data.apply(map_value, axis=1)
        # 5) Similarly, map chosen_label (0 or 1) to text if in predictor.prediction_map
        if chosen_label in predictor.prediction_map:
            def map_label(val):
                if val in [0, 1]:
                    return predictor.prediction_map[chosen_label][val]
                return f"Unknown label {val}"
            dist_data["LabelText"] = dist_data[chosen_label].apply(map_label)
        else:
            dist_data["LabelText"] = dist_data[chosen_label].astype(str)

        # 6) Now produce a bar chart with facet_col = FeatureName
        fig_distribution = px.bar(
            dist_data,
            x="FeatureValueText",
            y="count",
            color="LabelText", 
            facet_col="FeatureName",
            facet_col_wrap=4,     # how many facets per row
            title=f"Distribution of All Input Features vs. {chosen_label}",
            height=800
        )
        fig_distribution.update_layout(legend=dict(title=chosen_label))
        # (Optional) Adjust layout or text angle if you have many categories
        fig_distribution.update_xaxes(tickangle=45)
    else:
        # Fallback
        fig_distribution = px.bar(title=f"Label {chosen_label} not found in dataset. Distribution not available.")

    ############################################################################
    # C) Nearest Neighbors (Hamming Distance) with disclaimers & user-friendly text
    ############################################################################
    # "Nearest neighbor” methods for high-dimensional or purely categorical data can be non-trivial.
    # This demo uses a Hamming distance over all input features, picks K=5.
    # In real practice, you'd refine which features to use, how to encode them, etc.

    # 1) Build a DataFrame to compare with the user_input
    features_to_compare = [col for col in user_input_data if col in df.columns]
    user_series = user_input.iloc[0]

    # 2) Compute distances
    distances = []
    for idx, row in df[features_to_compare].iterrows():
        d = 0
        for col in features_to_compare:
            if row[col] != user_series[col]:
                d += 1
        distances.append(d)

    df_with_dist = df.copy()
    df_with_dist["distance"] = distances

    # 3) Sort and pick top K=5
    K = 5
    nearest_neighbors = df_with_dist.sort_values("distance", ascending=True).head(K)

    # 4) Show how many had the chosen_label=0 vs 1, but also map them
    #    We'll also demonstrate showing user-friendly text for each neighbor's feature values.
    #    However, if you have large K or many features, this can be big.
    if chosen_label in nearest_neighbors.columns:
        nn_label_0 = len(nearest_neighbors[nearest_neighbors[chosen_label] == 0])
        nn_label_1 = len(nearest_neighbors[nearest_neighbors[chosen_label] == 1])
        if chosen_label in predictor.prediction_map:
            label0_text = predictor.prediction_map[chosen_label][0]
            label1_text = predictor.prediction_map[chosen_label][1]
        else:
            label0_text = "Label=0"
            label1_text = "Label=1"
    else:
        nn_label_0 = nn_label_1 = 0
        label0_text = "Label=0"
        label1_text = "Label=1"

    # 5) Build an example table of those neighbors in user-friendly text
    neighbor_text_rows = []
    for idx, nn_row in nearest_neighbors.iterrows():
        # For each feature, map numeric -> user text
        row_str_parts = []
        row_str_parts.append(f"distance={nn_row['distance']}")
        for fcol in features_to_compare:
            val = nn_row[fcol]
            # try to map
            if fcol in reverse_mapping and val in reverse_mapping[fcol]:
                val_str = reverse_mapping[fcol][val]
            else:
                val_str = str(val)
            row_str_parts.append(f"{fcol}={val_str}")
        # For the label
        if chosen_label in nn_row:
            lbl_val = nn_row[chosen_label]
            if chosen_label in predictor.prediction_map and lbl_val in [0, 1]:
                lbl_str = predictor.prediction_map[chosen_label][lbl_val]
            else:
                lbl_str = str(lbl_val)
            row_str_parts.append(f"{chosen_label}={lbl_str}")
        neighbor_text_rows.append(" | ".join(row_str_parts))

    neighbor_text_block = "\n".join(neighbor_text_rows)

    similar_patient_markdown = (
        "### Nearest Neighbors (Simple Hamming Distance)\n"
        "“Nearest neighbor” methods for high-dimensional or purely categorical data can be non-trivial. "
        "This demo simply uses a Hamming distance over all input features and picks **K=5** neighbors.\n\n"
        "In a real application, you would refine which features are most relevant, how to encode them, "
        "and how many neighbors to select.\n\n"
        f"Among these **{K}** nearest neighbors:\n"
        f"- **{nn_label_0}** had {label0_text}\n"
        f"- **{nn_label_1}** had {label1_text}\n\n"
        "Below is a breakdown of each neighbor's key features in user-friendly text:\n\n"
        f"```\n{neighbor_text_block}\n```"
    )

    ############################################################################
    # Return 8 outputs
    ############################################################################
    return (
        formatted_results,              # 1) Prediction results (Textbox)
        severity,                       # 2) Mental Health Severity (Textbox)
        total_patient_count_markdown,   # 3) Total Patient Count (Markdown)
        fig_distribution,               # 4) Distribution Plot (Plot)
        similar_patient_markdown,       # 5) Nearest Neighbor Summary (Markdown)
        None,                           # 6) Placeholder if you need more plots
        None,                           # 7) Another placeholder
        None                            # 8) Another placeholder
    )

###############################################################################
# 6) Gradio Interface: We'll keep 8 outputs, but only use some in this demo
###############################################################################
def predict_with_text(
    YMDEYR, YMDERSUD5ANY, YMDEIMAD5YR, YMIMS5YANY, YMDELT, YMDEHARX, 
    YMDEHPRX, YMDETXRX, YMDEHPO, YMDEAUD5YR, YMIMI5YANY, YMIUD5YANY, 
    YMDESUD5ANYO, YNURSMDE, YSOCMDE, YCOUNMDE, YPSY1MDE, YPSY2MDE, 
    YHLTMDE, YDOCMDE, YTXMDEYR, YUSUITHKYR, YUSUIPLNYR, YUSUITHK, 
    YUSUIPLN, MDEIMPY, LVLDIFMEM2, YMSUD5YANY, YRXMDEYR
):
    # Validate that all required inputs are selected
    if not validate_inputs(
        YMDEYR, YMDERSUD5ANY, YMDEIMAD5YR, YMIMS5YANY, YMDELT, YMDEHARX, 
        YMDEHPRX, YMDETXRX, YMDEHPO, YMDEAUD5YR, YMIMI5YANY, YMIUD5YANY, 
        YMDESUD5ANYO, YNURSMDE, YSOCMDE, YCOUNMDE, YPSY1MDE, YPSY2MDE, 
        YHLTMDE, YDOCMDE, YTXMDEYR, YUSUITHKYR, YUSUIPLNYR, YUSUITHK, 
        YUSUIPLN, MDEIMPY, LVLDIFMEM2, YMSUD5YANY, YRXMDEYR
    ):
        return (
            "Please select all required fields.",  # Prediction Results
            "Validation Error",                    # Severity
            "No data",                             # Total Patient Count
            None,                                  # Distribution Plot
            "No data",                             # Nearest Neighbors
            None, None, None                       # Placeholders
        )
    
    # Map from user-friendly text to int
    user_inputs = {
        'YNURSMDE': input_mapping['YNURSMDE'][YNURSMDE],
        'YMDEYR': input_mapping['YMDEYR'][YMDEYR],
        'YSOCMDE': input_mapping['YSOCMDE'][YSOCMDE],
        'YMDESUD5ANYO': input_mapping['YMDESUD5ANYO'][YMDESUD5ANYO],
        'YMSUD5YANY': input_mapping['YMSUD5YANY'][YMSUD5YANY],
        'YUSUITHK': input_mapping['YUSUITHK'][YUSUITHK],
        'YMDETXRX': input_mapping['YMDETXRX'][YMDETXRX],
        'YUSUITHKYR': input_mapping['YUSUITHKYR'][YUSUITHKYR],
        'YMDERSUD5ANY': input_mapping['YMDERSUD5ANY'][YMDERSUD5ANY],
        'YUSUIPLNYR': input_mapping['YUSUIPLNYR'][YUSUIPLNYR],
        'YCOUNMDE': input_mapping['YCOUNMDE'][YCOUNMDE],
        'YPSY1MDE': input_mapping['YPSY1MDE'][YPSY1MDE],
        'YHLTMDE': input_mapping['YHLTMDE'][YHLTMDE],
        'YDOCMDE': input_mapping['YDOCMDE'][YDOCMDE],
        'YPSY2MDE': input_mapping['YPSY2MDE'][YPSY2MDE],
        'YMDEHARX': input_mapping['YMDEHARX'][YMDEHARX],
        'LVLDIFMEM2': input_mapping['LVLDIFMEM2'][LVLDIFMEM2],
        'MDEIMPY': input_mapping['MDEIMPY'][MDEIMPY],
        'YMDEHPO': input_mapping['YMDEHPO'][YMDEHPO],
        'YMIMS5YANY': input_mapping['YMIMS5YANY'][YMIMS5YANY],
        'YMDEIMAD5YR': input_mapping['YMDEIMAD5YR'][YMDEIMAD5YR],
        'YMIUD5YANY': input_mapping['YMIUD5YANY'][YMIUD5YANY],
        'YMDEHPRX': input_mapping['YMDEHPRX'][YMDEHPRX],
        'YMIMI5YANY': input_mapping['YMIMI5YANY'][YMIMI5YANY],
        'YUSUIPLN': input_mapping['YUSUIPLN'][YUSUIPLN],
        'YTXMDEYR': input_mapping['YTXMDEYR'][YTXMDEYR],
        'YMDEAUD5YR': input_mapping['YMDEAUD5YR'][YMDEAUD5YR],
        'YRXMDEYR': input_mapping['YRXMDEYR'][YRXMDEYR],
        'YMDELT': input_mapping['YMDELT'][YMDELT]
    }
    
    # Pass our mapped values into the original 'predict' function
    return predict(**user_inputs)

###############################################################################
# 7) Define and Launch Gradio Interface
###############################################################################
import sys

# We have 8 outputs (some are placeholders)
outputs = [
    gr.Textbox(label="Prediction Results", lines=30),
    gr.Textbox(label="Mental Health Severity", lines=4),
    gr.Markdown(label="Total Patient Count"),
    gr.Plot(label="Distribution of All Input Features vs. One Label"),
    gr.Markdown(label="Nearest Neighbors Summary"),
    gr.Plot(label="Placeholder Plot"),
    gr.Plot(label="Placeholder Plot"),
    gr.Plot(label="Placeholder Plot")
]

# Define the inputs
inputs = [
    # Major Depressive Episode (MDE) questions
    gr.Dropdown(list(input_mapping['YMDEYR'].keys()), label="YMDEYR: PAST YEAR MDE?"),
    gr.Dropdown(list(input_mapping['YMDERSUD5ANY'].keys()), label="YMDERSUD5ANY: MDE OR SUBSTANCE USE DISORDER - ANY"),
    gr.Dropdown(list(input_mapping['YMDEIMAD5YR'].keys()), label="YMDEIMAD5YR: MDE + ALCOHOL USE DISORDER?"),
    gr.Dropdown(list(input_mapping['YMIMS5YANY'].keys()), label="YMIMS5YANY: MDE + SUBSTANCE USE DISORDER?"),
    gr.Dropdown(list(input_mapping['YMDELT'].keys()), label="YMDELT: EVER HAD MDE LIFETIME?"),
    gr.Dropdown(list(input_mapping['YMDEHARX'].keys()), label="YMDEHARX: SAW HEALTH PROF + MEDS FOR MDE"),
    gr.Dropdown(list(input_mapping['YMDEHPRX'].keys()), label="YMDEHPRX: SAW HEALTH PROF OR MEDS FOR MDE"),
    gr.Dropdown(list(input_mapping['YMDETXRX'].keys()), label="YMDETXRX: TREATMENT/COUNSELING FOR MDE"),
    gr.Dropdown(list(input_mapping['YMDEHPO'].keys()), label="YMDEHPO: HEALTH PROF ONLY FOR MDE"),
    gr.Dropdown(list(input_mapping['YMDEAUD5YR'].keys()), label="YMDEAUD5YR: MDE + ALCOHOL USE DISORDER"),
    gr.Dropdown(list(input_mapping['YMIMI5YANY'].keys()), label="YMIMI5YANY: MDE + ILL DRUG USE DISORDER"),
    gr.Dropdown(list(input_mapping['YMIUD5YANY'].keys()), label="YMIUD5YANY: MDE + ILL DRUG USE DISORDER"),
    gr.Dropdown(list(input_mapping['YMDESUD5ANYO'].keys()), label="YMDESUD5ANYO: MDE vs. SUD vs. BOTH vs. NEITHER"),
    
    # Consultations
    gr.Dropdown(list(input_mapping['YNURSMDE'].keys()), label="YNURSMDE: NURSE / OT FOR MDE"),
    gr.Dropdown(list(input_mapping['YSOCMDE'].keys()), label="YSOCMDE: SOCIAL WORKER FOR MDE"),
    gr.Dropdown(list(input_mapping['YCOUNMDE'].keys()), label="YCOUNMDE: COUNSELOR FOR MDE"),
    gr.Dropdown(list(input_mapping['YPSY1MDE'].keys()), label="YPSY1MDE: PSYCHOLOGIST FOR MDE"),
    gr.Dropdown(list(input_mapping['YPSY2MDE'].keys()), label="YPSY2MDE: PSYCHIATRIST FOR MDE"),
    gr.Dropdown(list(input_mapping['YHLTMDE'].keys()), label="YHLTMDE: HEALTH PROF FOR MDE"),
    gr.Dropdown(list(input_mapping['YDOCMDE'].keys()), label="YDOCMDE: GP/FAMILY MD FOR MDE"),
    gr.Dropdown(list(input_mapping['YTXMDEYR'].keys()), label="YTXMDEYR: DOCTOR/HEALTH PROF FOR MDE THIS YEAR"),
    
    # Suicidal thoughts / plans
    gr.Dropdown(list(input_mapping['YUSUITHKYR'].keys()), label="YUSUITHKYR: SERIOUSLY THOUGHT ABOUT KILLING SELF"),
    gr.Dropdown(list(input_mapping['YUSUIPLNYR'].keys()), label="YUSUIPLNYR: MADE PLANS TO KILL SELF"),
    gr.Dropdown(list(input_mapping['YUSUITHK'].keys()), label="YUSUITHK: THINK ABOUT KILLING SELF (12 MONTHS)"),
    gr.Dropdown(list(input_mapping['YUSUIPLN'].keys()), label="YUSUIPLN: MADE PLANS TO KILL SELF (12 MONTHS)"),

    # Impairment
    gr.Dropdown(list(input_mapping['MDEIMPY'].keys()), label="MDEIMPY: MDE WITH SEVERE ROLE IMPAIRMENT?"),
    gr.Dropdown(list(input_mapping['LVLDIFMEM2'].keys()), label="LVLDIFMEM2: DIFFICULTY REMEMBERING/CONCENTRATING"),
    gr.Dropdown(list(input_mapping['YMSUD5YANY'].keys()), label="YMSUD5YANY: MDE + SUBSTANCE USE DISORDER?"),
    gr.Dropdown(list(input_mapping['YRXMDEYR'].keys()), label="YRXMDEYR: USED MEDS FOR MDE IN PAST YEAR?")
]

# Custom CSS (optional)
custom_css = """
    .gradio-container * {
        color: #1B1212 !important;
    }
    .gradio-container .form .form-group label {
        color: #1B1212 !important;
    }
    .gradio-container .output-textbox, 
    .gradio-container .output-textbox textarea {
        color: #1B1212 !important;
    }
    .gradio-container .label,
    .gradio-container .input-label {
        color: #1B1212 !important;
    }
"""

# Build the interface
interface = gr.Interface(
    fn=predict_with_text,
    inputs=inputs,
    outputs=outputs,
    title="Adolescents with Substance Use Mental Health Screening (NSDUH Data)",
    css=custom_css,
)

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