import gradio as gr
import pandas as pd
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
import shap
import numpy as np
from io import BytesIO

# -------------------------------
# Load Model
# -------------------------------
model = joblib.load("model/performance_pipeline.pkl")

categorical_features = ['school','sex','address','famsize','Pstatus','Mjob','Fjob','reason','guardian',
                        'schoolsup','famsup','paid','activities','nursery','higher','internet','romantic','dataset']
numeric_features = ['age','Medu','Fedu','traveltime','studytime','failures','famrel','freetime',
                    'goout','Dalc','Walc','health','absences','G1','G2','G3']

# -------------------------------
# Single Prediction Function
# -------------------------------
def single_prediction(*inputs):
    # Map inputs to dataframe
    data = dict(zip(categorical_features + numeric_features, inputs))
    df = pd.DataFrame([data])

    if 'dataset' not in df.columns:
        df['dataset'] = 'student_mat'

    pred = model.predict(df)[0]

    return f"Predicted Performance: {pred}"

# -------------------------------
# Batch Prediction Function
# -------------------------------
def batch_prediction(file):
    df = pd.read_csv(file.name)

    if 'dataset' not in df.columns:
        df['dataset'] = 'student_mat'

    preds = model.predict(df)
    df["Prediction"] = preds

    # Plot counts
    pred_counts = df["Prediction"].value_counts()
    fig, ax = plt.subplots()
    sns.barplot(x=pred_counts.index, y=pred_counts.values, palette="coolwarm", ax=ax)
    ax.set_ylabel("Count")

    # Save plot to buffer
    buf = BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)

    return df.head(), buf

# -------------------------------
# Gradio Interfaces
# -------------------------------
# Single prediction UI
single_inputs = []
for col in categorical_features:
    if col == 'school':
        single_inputs.append(gr.Dropdown(["GP", "MS"], label=col))
    elif col == 'address':
        single_inputs.append(gr.Dropdown(["U", "R"], label=col))
    elif col == 'famsize':
        single_inputs.append(gr.Dropdown(["GT3", "LE3"], label=col))
    elif col == 'Pstatus':
        single_inputs.append(gr.Dropdown(["T", "A"], label=col))
    elif col in ['Mjob','Fjob']:
        single_inputs.append(gr.Dropdown(["teacher","health","services","at_home","other"], label=col))
    elif col == 'reason':
        single_inputs.append(gr.Dropdown(["home","reputation","course","other"], label=col))
    elif col == 'guardian':
        single_inputs.append(gr.Dropdown(["mother","father","other"], label=col))
    elif col in ['schoolsup','famsup','paid','activities','nursery','higher','internet','romantic']:
        single_inputs.append(gr.Dropdown(["yes", "no"], label=col))
    else:
        single_inputs.append(gr.Textbox(label=col))

for col in numeric_features:
    single_inputs.append(gr.Number(label=col))

single_demo = gr.Interface(
    fn=single_prediction,
    inputs=single_inputs,
    outputs="text",
    title="🎓 StudentPass - Single Prediction"
)

# Batch prediction UI
batch_demo = gr.Interface(
    fn=batch_prediction,
    inputs=gr.File(label="Upload CSV"),
    outputs=[gr.Dataframe(), gr.Image(type="pil")],
    title="🎓 StudentPass - Batch Prediction"
)

# Combine into Tabs
demo = gr.TabbedInterface([single_demo, batch_demo], ["Single Prediction", "Batch Prediction"])

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