import streamlit as st
import joblib
import pandas as pd
import numpy as np

# --- Configuration ---
MODEL_PATH = 'src/kn_perormance_prediction.joblib'
ENCODERS_PATH = 'src/label_encoders.joblib'

# REAL features the model expects (first 10 columns of df during training)
FEATURES = [
    'gender', 
    'race/ethnicity', 
    'parental level of education', 
    'lunch', 
    'test preparation course', 
    'math score', 
    'reading score', 
    'writing score',
    'total_score',
    'percentage'
]

RACE_MAP = {
    'group A': 1, 'group B': 2, 'group C': 3,
    'group D': 4, 'group E': 5
}

@st.cache_resource
def load_assets():
    try:
        model = joblib.load(MODEL_PATH)
        encoders = joblib.load(ENCODERS_PATH)
        return model, encoders
    except Exception as e:
        st.error(f"Error loading model or encoders: {e}")
        return None, None


def preprocess_and_predict(model, encoders, input_data):

    df_input = pd.DataFrame([input_data])

    # 1. Map race/ethnicity exactly as in training
    df_input['race/ethnicity'] = df_input['race/ethnicity'].map(RACE_MAP)

    # 2. Apply Label Encoding for categorical inputs
    for col, le in encoders.items():
        df_input[col] = le.transform(df_input[col])

    # 3. Compute total_score and percentage (missing features)
    df_input['total_score'] = (
        df_input['math score'] +
        df_input['reading score'] +
        df_input['writing score']
    )

    df_input['percentage'] = df_input['total_score'] / 3

    # 4. Final array in the correct feature order
    final_input_array = df_input[FEATURES].values

    # 5. Predict grade
    prediction_label = model.predict(final_input_array)

    return prediction_label[0]


# --- Streamlit UI ---
st.set_page_config(page_title="Student Performance", layout="centered")
st.title("📚 Student Final Grade Predictor (A–E)")
st.markdown("Enter student characteristics and scores to predict the final letter grade.")

model, encoders = load_assets()

if model is not None and encoders is not None:
    st.sidebar.header("Student Information")

    gender = st.sidebar.selectbox("Gender:", options=['male', 'female'])
    race = st.sidebar.selectbox("Race/Ethnicity:", options=list(RACE_MAP.keys()))
    parental_education = st.sidebar.selectbox(
        "Parental Education:",
        options=[
            'some high school', 'high school', 'some college',
            "associate's degree", "bachelor's degree", "master's degree"
        ]
    )
    lunch = st.sidebar.selectbox("Lunch:", options=['standard', 'free/reduced'])
    prep_course = st.sidebar.selectbox("Test Prep Course:", options=['none', 'completed'])

    st.sidebar.header("Scores (0-100)")
    math = st.sidebar.slider("Math Score:", 0, 100, 70)
    reading = st.sidebar.slider("Reading Score:", 0, 100, 75)
    writing = st.sidebar.slider("Writing Score:", 0, 100, 72)

    input_data = {
        'gender': gender,
        'race/ethnicity': race,
        'parental level of education': parental_education,
        'lunch': lunch,
        'test preparation course': prep_course,
        'math score': math,
        'reading score': reading,
        'writing score': writing
    }

    if st.button("Predict Final Grade"):
        with st.spinner("Calculating prediction..."):
            predicted_grade = preprocess_and_predict(model, encoders, input_data)
            
            st.success("Prediction Successful!")
            st.markdown("### Predicted Letter Grade:")
            st.markdown(f"**{predicted_grade}**")