import gradio as gr
import pickle
import numpy as np
import pandas as pd

class ModelPredictor:
    def __init__(self, model_path, model_filenames):
        self.model_path = model_path
        self.model_filenames = model_filenames
        self.models = self.load_models()
        self.prediction_map = {
            'YOWRDISC': ["Did not feel discouraged about life every day", "Felt discouraged about life every day"],
            'YO_MDEA3': ["Does not have changes in appetite or weight", "Has changes in appetite or weight"],
            'YOPB2WK': ["Did not feel this way for 2 weeks or more with other problems during the past 12 months", "Felt this way for 2 weeks or more with other problems during the past 12 months"],
            'YOWRDCSN': ["Able to make decisions when problems were worst", "Unable to make decisions when problems were worst"],
            'YODPDISC': ["Did not feel discouraged about life when sad/empty/depressed", "Felt discouraged about life when sad/empty/depressed"],
            'YOWRELES': ["Did not eat less than usual every day when problems were worst", "Ate less than usual every day when problems were worst"],
            'YORX12MO': ["Not currently receiving treatment/counseling for mood", "Currently receiving treatment/counseling for mood"],
            'YODPLSIN': ["Did not lose interest in things when sad/empty/depressed", "Lost interest in things when sad/empty/depressed"],
            'YODSCEV': ["Never had a period of time where felt discouraged or hopeless", "Ever had a period of time where felt discouraged or hopeless"],
            'YO_MDEA8': ["Ability to concentrate or make decisions", "Inability to concentrate or make decisions"],
            'YOWRHRS': ["Time that most severe/frequent mood lasted less than 1 hour", "Time that most severe/frequent mood lasted at least 1 hour but no more than 3 hours", "Time that most severe/frequent mood lasted at least 3 hours but no more than 5 hours"],
            'YOWRIMP': ["Mood so severe couldn't do daily activities not very often", "Mood so severe couldn't do daily activities sometimes", "Mood so severe couldn't do daily activities often"],
            'YO_MDEA2': ["Did not lose interest or pleasure in most things", "Lost interest or pleasure in most things"],
            'YODPPROB': ["Never had other problems during 2 weeks or longer", "Ever had other problems during 2 weeks or longer"],
            'YOWRPROB': ["No particular time that is the worst one ever", "One particular time that is the worst one ever"],
            'YOLOSEV': ["Never had a period of time lost interest in enjoyable things", "Ever had a period of time lost interest in enjoyable things"],
            'YO_MDEA5': ["Others did not notice that was restless or lethargic", "Others noticed that was restless or lethargic"],
            'YODSMMDE': ["Has less than 5 symptoms", "Has 5 or more symptoms"],
            'YODPR2WK': ["No time when feelings lasted every day for 2 weeks or longer", "Time when feelings lasted 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):
        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):
        combined_predictions = np.concatenate(predictions)
        majority_vote = np.bincount(combined_predictions).argmax()
        return majority_vote

    def evaluate_severity(self, majority_vote_count):
        if majority_vote_count > 10:
            return "Mental health severity: Severe"
        elif majority_vote_count > 7:
            return "Mental health severity: Moderate"
        elif majority_vote_count > 4:
            return "Mental health severity: Low"
        else:
            return "Mental health severity: Very Low"

model_filenames = [
    "YOWRDISC.pkl", "YO_MDEA3.pkl", "YOPB2WK.pkl", "YOWRDCSN.pkl", "YODPDISC.pkl", 
    "YOWRELES.pkl", "YORX12MO.pkl", "YODPLSIN.pkl", "YOWRIMP.pkl", "YODPPROB.pkl", 
    "YOWRPROB.pkl", "YOLOSEV.pkl", "YO_MDEA5.pkl", "YODSMMDE.pkl", "YODPR2WK.pkl"
]
model_path = "models/"
predictor = ModelPredictor(model_path, model_filenames)

def predict(YNURSMDE, YMDEYR, YSOCMDE, YUCOSUITHK, YMDESUD5ANYO, YMSUD5YANY, YUSUITHK, YMDETXRX, YUSUITHKYR, YMDERSUD5ANY, YUSUIPLNYR, YCOUNMDE, YPSY1MDE, YHLTMDE, YDOCMDE, YPSY2MDE, YMDEHARX, LVLDIFMEM2, MDEIMPY, YMDEHPO, CONEGMH, COCLNEGMH, YMIMS5YANY, YMDEIMAD5YR, YMIUD5YANY, YUCOSUIPLN, YMDEHPRX, YMIMI5YANY, YUSUIPLN, YTXMDEYR, YMDEAUD5YR, YRXMDEYR, YMDELT):
    user_input_data = {
        'YNURSMDE': [int(YNURSMDE)],
        'YMDEYR': [int(YMDEYR)],
        'YSOCMDE': [int(YSOCMDE)],
        'YUCOSUITHK': [int(YUCOSUITHK)],
        '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)],
        'CONEGMH': [int(CONEGMH)],
        'COCLNEGMH': [int(COCLNEGMH)],
        'YMIMS5YANY': [int(YMIMS5YANY)],
        'YMDEIMAD5YR': [int(YMDEIMAD5YR)],
        'YMIUD5YANY': [int(YMIUD5YANY)],
        'YUCOSUIPLN': [int(YUCOSUIPLN)],
        '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)
    predictions = predictor.make_predictions(user_input)
    majority_vote = predictor.get_majority_vote(predictions)
    majority_vote_count = sum([1 for pred in np.concatenate(predictions) if pred == 1])
    severity = predictor.evaluate_severity(majority_vote_count)

    results = []
    for i, pred in enumerate(predictions):
        model_name = model_filenames[i].split('.')[0]
        pred_value = pred[0]
        if model_name in predictor.prediction_map:
            if pred_value < len(predictor.prediction_map[model_name]):
                results.append(f"Prediction from model {model_name}: {predictor.prediction_map[model_name][pred_value]}")
            else:
                results.append(f"Prediction from model {model_name}: Unknown prediction value {pred_value}")
        else:
            results.append(f"Prediction from model {model_name}: Unknown model")

    return results, severity

inputs = [
    gr.Dropdown(["1", "0"], label="YNURSMDE (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "2"], label="YMDEYR (1 = Yes, 2 = No)"),
    gr.Dropdown(["1", "0"], label="YSOCMDE (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "2"], label="YUCOSUITHK (1 = Yes, 2 = No)"),
    gr.Dropdown(["1", "2", "3", "4"], label="YMDESUD5ANYO (1 = SUD only, no MDE, 2 = MDE only, no SUD, 3 = SUD and MDE, 4 = Neither SUD or MDE)"),
    gr.Dropdown(["1", "0"], label="YMSUD5YANY (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "2", "3", "4"], label="YUSUITHK (1 = Yes, 2 = No, 3 = I'm not sure, 4 = I don't want to answer)"),
    gr.Dropdown(["1", "0"], label="YMDETXRX (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "2", "3", "4"], label="YUSUITHKYR (1 = Yes, 2 = No, 3 = I'm not sure, 4 = I don't want to answer)"),
    gr.Dropdown(["1", "0"], label="YMDERSUD5ANY (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "2", "3", "4"], label="YUSUIPLNYR (1 = Yes, 2 = No, 3 = I'm not sure, 4 = I don't want to answer)"),
    gr.Dropdown(["1", "0"], label="YCOUNMDE (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "0"], label="YPSY1MDE (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "0"], label="YHLTMDE (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "0"], label="YDOCMDE (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "0"], label="YPSY2MDE (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "0"], label="YMDEHARX (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "2", "3"], label="LVLDIFMEM2 (1 = No Difficulty, 2 = Some difficulty, 3 = A lot of difficulty or cannot do at all)"),
    gr.Dropdown(["1", "2"], label="MDEIMPY (1 = Yes, 2 = No)"),
    gr.Dropdown(["1", "0"], label="YMDEHPO (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "2", "3", "4", "5"], label="CONEGMH (1 = Not at all, 2 = A little, 3 = Some, 4 = Quite a bit, 5 = A lot)"),
    gr.Dropdown(["1", "2", "3"], label="COCLNEGMH (1 = Not at all, 2 = A little or some, 3 = Quite a bit or a lot)"),
    gr.Dropdown(["1", "0"], label="YMIMS5YANY (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "0"], label="YMDEIMAD5YR (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "0"], label="YMIUD5YANY (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "2"], label="YUCOSUIPLN (1 = Yes, 2 = No)"),
    gr.Dropdown(["1", "0"], label="YMDEHPRX (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "0"], label="YMIMI5YANY (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "2", "3", "4"], label="YUSUIPLN (1 = Yes, 2 = No, 3 = I'm not sure, 4 = I don't want to answer)"),
    gr.Dropdown(["1", "0"], label="YTXMDEYR (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "0"], label="YMDEAUD5YR (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "0"], label="YRXMDEYR (1 = Yes, 0 = No)"),
    gr.Dropdown(["1", "2"], label="YMDELT (1 = Yes, 2 = No)")
]

outputs = [
    gr.Textbox(label="Prediction Results"),
    gr.Textbox(label="Mental Health Severity")
]

gr.Interface(fn=predict, inputs=inputs, outputs=outputs, title="Adolescent Mental Health Prediction").launch()