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

# Load trained models
with open('rf_hacathon_fullstk.pkl', 'rb') as f1:
    rf_fullstk = pickle.load(f1)
with open('rf_hacathon_prodengg.pkl', 'rb') as f2:
    rf_prodengg = pickle.load(f2)
with open('rf_hacathon_mkt.pkl', 'rb') as f3:
    rf_mkt = pickle.load(f3)

# Define prediction function
def predict_placed(degree_p, internship, DSA, java, management, leadership, communication, sales, model_name):
    if model_name == 'Full Stack':
        new_data = pd.DataFrame({
            'degree_p': degree_p,
            'internship': internship,
            'DSA': DSA,
            'java': java,
            'management': 0,
            'leadership': 0,
            'communication': 0,
            'sales': 0
        }, index=[0])
        model = rf_fullstk
    elif model_name == 'Product Engineering':
        new_data = pd.DataFrame({
            'degree_p': degree_p,
            'internship': internship,
            'DSA': 0,
            'java': 0,
            'management': management,
            'leadership': leadership,
            'communication': 0,
            'sales': 0
        }, index=[0])
        model = rf_prodengg
    elif model_name == 'Marketing':
        new_data = pd.DataFrame({
            'degree_p': degree_p,
            'internship': internship,
            'DSA': 0,
            'java': 0,
            'management': 0,
            'leadership': 0,
            'communication': communication,
            'sales': sales
        }, index=[0])
        model = rf_mkt

    prediction = model.predict(new_data)
    probability = model.predict_proba(new_data)[0][1]

    if prediction == 1:
        result = 'Placed'
        probability_message = f"You will be placed with a probability of {probability:.2f}"
    else:
        result = 'Not Placed'
        probability_message = ""

    return result, probability_message

# Create Gradio interface
inputs = [
    gr.inputs.Number(label='Degree Percentage'),
    gr.inputs.Radio(label='Internship', choices=[0, 1]),
    gr.inputs.Radio(label='Data Structures & Algorithms', choices=[0, 1]),
    gr.inputs.Radio(label='Java', choices=[0, 1]),
    gr.inputs.Radio(label='Management Skills', choices=[0, 1]),
    gr.inputs.Radio(label='Leadership Skills', choices=[0, 1]),
    gr.inputs.Radio(label='Communication Skills', choices=[0, 1]),
    gr.inputs.Radio(label='Sales Skills', choices=[0, 1]),
    gr.inputs.Dropdown(label='Model Name', choices=['Full Stack', 'Product Engineering', 'Marketing'])
]

outputs = [
    gr.outputs.Textbox(label='Placement Result'),
    gr.outputs.Textbox(label='Placement Probability')
]

app = gr.Interface(
    fn=predict_placed,
    inputs=inputs,
    outputs=outputs,
    title='Placement Prediction',
    description='Predict placement outcome based on given inputs',
    allow_flagging=False
)

# Run the app
if __name__ == '__main__':
    app.run()