import gradio as gr
import joblib
import numpy as np

# Load the saved models and pipelines
anxiety_model = joblib.load('Anxiety_best_model.pkl')
anxiety_pipeline = joblib.load('Anxiety_best_pipeline.pkl')

depression_model = joblib.load('Depression_best_model.pkl')
depression_pipeline = joblib.load('Depression_best_pipeline.pkl')

insomnia_model = joblib.load('Insomnia_best_model.pkl')
insomnia_pipeline = joblib.load('Insomnia_best_pipeline.pkl')

ocd_model = joblib.load('OCD_best_model.pkl')
ocd_pipeline = joblib.load('OCD_best_pipeline.pkl')

# Define the prediction functions
def Anxiety_predict(input1, input2, input3, input4, input5, input6):
    inputs = np.array([input1, input2, input3, input4, input5, input6]).reshape(1, -1)
    inputs_scaled = anxiety_pipeline.transform(inputs)
    prediction = anxiety_model.predict(inputs_scaled)
    return prediction[0], f"The predicted value is {prediction[0]:.2f}"

def Depression_predict(input1, input2, input3, input4, input5, input6):
    inputs = np.array([input1, input2, input3, input4, input5, input6]).reshape(1, -1)
    inputs_scaled = depression_pipeline.transform(inputs)
    prediction = depression_model.predict(inputs_scaled)
    return prediction[0], f"The predicted value is {prediction[0]:.2f}"

def Insomnia_predict(input1, input2, input3, input4, input5, input6):
    inputs = np.array([input1, input2, input3, input4, input5, input6]).reshape(1, -1)
    inputs_scaled = insomnia_pipeline.transform(inputs)
    prediction = insomnia_model.predict(inputs_scaled)
    return prediction[0], f"The predicted value is {prediction[0]:.2f}"

def OCD_predict(input1, input2, input3, input4, input5, input6):
    inputs = np.array([input1, input2, input3, input4, input5, input6]).reshape(1, -1)
    inputs_scaled = ocd_pipeline.transform(inputs)
    prediction = ocd_model.predict(inputs_scaled)
    return prediction[0], f"The predicted value is {prediction[0]:.2f}"

# Define the Gradio interfaces
anxiety_iface = gr.Interface(
    fn=Anxiety_predict,
    inputs=[
        gr.Number(label="Age"),
        gr.Number(label="Hours per day"),
        gr.Number(label="Depression"),
        gr.Number(label="Insomnia"),
        gr.Number(label="OCD"),
        gr.Number(label="BPM")
    ],
    outputs=[
        gr.Number(label="Predicted Value"),
        gr.Textbox(label="Prediction Description")
    ],
    title="Music & Mental Health Predictor - Anxiety",
    description="Enter 5 numeric values to get a prediction.",
    theme=gr.themes.Soft(),
    examples=[[18, 3, 0, 1, 0, 156]]
)
depression_iface = gr.Interface(
    fn=Depression_predict,
    inputs=[
        gr.Number(label="Age"),
        gr.Number(label="Hours per day"),
        gr.Number(label="Insomnia"),
        gr.Number(label="Anxiety"),
        gr.Number(label="OCD"),
        gr.Number(label="BPM")
    ],
    outputs=[
        gr.Number(label="Predicted Value"),
        gr.Textbox(label="Prediction Description")
    ],
    title="Music & Mental Health Predictor - Depression",
    description="Enter 5 numeric values to get a prediction.",
    theme=gr.themes.Soft(),
    examples=[[18, 3, 0, 1, 0, 156]]
)
insomnia_iface = gr.Interface(
    fn=Insomnia_predict,
    inputs=[
        gr.Number(label="Age"),
        gr.Number(label="Hours per day"),
        gr.Number(label="Depression"),
        gr.Number(label="Anxiety"),
        gr.Number(label="OCD"),
        gr.Number(label="BPM")
    ],
    outputs=[
        gr.Number(label="Predicted Value"),
        gr.Textbox(label="Prediction Description")
    ],
    title="Music & Mental Health Predictor - Insomnia",
    description="Enter 5 numeric values to get a prediction.",
    theme=gr.themes.Soft(),
    examples=[[18, 3, 0, 1, 0, 156]]
)

ocd_iface = gr.Interface(
    fn=OCD_predict,
    inputs=[
        gr.Number(label="Age"),
        gr.Number(label="Hours per day"),
        gr.Number(label="Insomnia"),
        gr.Number(label="Anxiety"),
        gr.Number(label="Depression"),
        gr.Number(label="BPM")
    ],
    outputs=[
        gr.Number(label="Predicted Value"),
        gr.Textbox(label="Prediction Description")
    ],
    title="Music & Mental Health Predictor - OCD",
    description="Enter 5 numeric values to get a prediction.",
    theme=gr.themes.Soft(),
    examples=[[18, 3, 0, 1, 0, 156]]
)

# Create a tabbed interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    with gr.Tab("Anxiety"):
        anxiety_iface.render()
    with gr.Tab("Depression"):
        depression_iface.render()
    with gr.Tab("Insomnia"):
        insomnia_iface.render()
    with gr.Tab("OCD"):
        ocd_iface.render()

# Launch the interface
demo.launch(debug=True)