import streamlit as st
import numpy as np
import joblib

# ===== Load models =====
model_full = joblib.load("random_forest_model_full.pkl")
model_important = joblib.load("random_forest_model_importainfeature.pkl")

# ===== Setup page =====
st.set_page_config(page_title="Dự đoán tốt nghiệp đúng hạn", page_icon="🎓", layout="wide")

# ===== Custom CSS =====
st.markdown("""
    <style>
    html, body, [class*="css"]  {
        font-family: 'Poppins', sans-serif;
    }
    .block-container {
        padding: 2rem 2rem 2rem 2rem;
    }
    .stButton button {
        background: linear-gradient(90deg, #4F8BF9 0%, #8E2DE2 100%);
        border: none;
        padding: 0.75em 1.5em;
        border-radius: 8px;
        color: white;
        font-size: 18px;
        font-weight: bold;
    }
    </style>
""", unsafe_allow_html=True)

# ===== HEADER =====
st.markdown("<h1 style='text-align: center; color: #003366; font-size: 28px;'>🎓 DỰ ĐOÁN KHẢ NĂNG TỐT NGHIỆP ĐÚNG HẠN</h1>", unsafe_allow_html=True)
st.markdown("<h4 style='text-align: center; color: #666;'>Áp dụng cho sinh viên năm 3</h4>", unsafe_allow_html=True)
st.write("---")

# ===========================
#     BỘ 8 VÍ DỤ
# ===========================
sample_options = [
    "Không ví dụ",
    "CNTT (Full) – Đúng hạn",
    "CNTT (Full) – Không đúng hạn",
    "Kinh tế (Full) – Đúng hạn",
    "Kinh tế (Full) – Không đúng hạn",
    "CNTT (Simple) – Đúng hạn",
    "CNTT (Simple) – Không đúng hạn",
    "Kinh tế (Simple) – Đúng hạn",
    "Kinh tế (Simple) – Không đúng hạn",
]

sample_option = st.selectbox("📝 Chọn ví dụ:", sample_options)

# ===== FULL DATA EXAMPLES =====
sample_cntt_full_ok = {
    "name": "Nguyễn Văn A",
    "student_id": "10117001",
    "major": "Công nghệ thông tin",
    "semester_data": [
        0,0,0,0,17,7,7,8.5,6,
        0,0,0,0,17,7,7,8.3,5,
        0,0,0,0,17,7,7,8.0,5,
        0,0,0,0,17,7,6,8.2,6,
        0,0,0,0,19,7,6,8.1,6,
        0,0,3,1,19,7,6,7.9,5
    ]
}

sample_cntt_full_fail = {
    "name": "Nguyễn Văn B",
    "student_id": "10117002",
    "major": "Công nghệ thông tin",
    "semester_data": [
        1,3,5,2,15,6,4,5.5,2,
        2,3,6,3,16,6,3,5.2,2,
        2,4,7,3,17,7,3,4.9,1,
        3,4,8,4,17,7,3,4.8,1,
        3,5,9,4,18,7,3,4.7,1,
        4,6,10,5,18,7,2,4.5,0
    ]
}

sample_kinhte_full_ok = {
    "name": "Trần Thị C",
    "student_id": "11418001",
    "major": "Kinh tế",
    "semester_data": [
        0,0,0,0,16,6,6,7.0,5,
        0,0,0,0,18,7,7,8.2,6,
        0,0,0,0,17,7,7,7.8,5,
        0,0,0,0,17,7,7,8.0,6,
        0,0,0,0,19,7,7,8.3,6,
        0,0,0,0,19,7,7,8.1,6
    ]
}

sample_kinhte_full_fail = {
    "name": "Trần Thị D",
    "student_id": "11418002",
    "major": "Kinh tế",
    "semester_data": [
        1,2,3,2,15,6,4,5.6,2,
        1,3,4,3,15,6,4,5.3,2,
        2,4,5,3,16,7,4,5.2,2,
        2,5,6,3,17,7,4,5.0,1,
        3,6,7,4,17,7,3,4.9,1,
        4,7,8,5,18,7,3,4.7,1
    ]
}

# ===== SIMPLE DATA EXAMPLES =====
sample_cntt_simple_ok = {
    "name": "Nguyễn Văn E",
    "student_id": "10118001",
    "major": "Công nghệ thông tin",
    "semester_data": [
        17,0,7.2,
        17,0,7.0,
        18,0,6.8,
        18,0,6.9,
        19,0,7.1,
        19,0,7.0
    ]
}

sample_cntt_simple_fail = {
    "name": "Nguyễn Văn F",
    "student_id": "10118002",
    "major": "Công nghệ thông tin",
    "semester_data": [
        12,10,5.2,
        12,12,5.0,
        13,14,4.8,
        14,13,4.7,
        15,12,4.9,
        15,13,4.6
    ]
}

sample_kinhte_simple_ok = {
    "name": "Trần Thị G",
    "student_id": "11418003",
    "major": "Kinh tế",
    "semester_data": [
        17,0,7.5,
        17,0,7.2,
        17,0,7.0,
        18,0,7.3,
        18,0,7.4,
        19,0,7.5
    ]
}

sample_kinhte_simple_fail = {
    "name": "Trần Thị H",
    "student_id": "11418004",
    "major": "Kinh tế",
    "semester_data": [
        12,10,5.3,
        12,11,5.0,
        13,12,4.9,
        13,13,4.8,
        14,12,4.7,
        14,13,4.5
    ]
}

# ==============================
#   LOAD CORRESPONDING EXAMPLE
# ==============================
example_data = None
if sample_option == "CNTT (Full) – Đúng hạn":
    example_data = sample_cntt_full_ok
elif sample_option == "CNTT (Full) – Không đúng hạn":
    example_data = sample_cntt_full_fail
elif sample_option == "Kinh tế (Full) – Đúng hạn":
    example_data = sample_kinhte_full_ok
elif sample_option == "Kinh tế (Full) – Không đúng hạn":
    example_data = sample_kinhte_full_fail
elif sample_option == "CNTT (Simple) – Đúng hạn":
    example_data = sample_cntt_simple_ok
elif sample_option == "CNTT (Simple) – Không đúng hạn":
    example_data = sample_cntt_simple_fail
elif sample_option == "Kinh tế (Simple) – Đúng hạn":
    example_data = sample_kinhte_simple_ok
elif sample_option == "Kinh tế (Simple) – Không đúng hạn":
    example_data = sample_kinhte_simple_fail

# ===============================
#     PERSONAL INFO
# ===============================
if example_data:
    name = st.text_input("👤 Họ và tên", value=example_data["name"])
    student_id = st.text_input("🎓 Mã sinh viên", value=example_data["student_id"])
    major = st.selectbox("📚 Ngành học",
                         ["Công nghệ thông tin", "Kinh tế"],
                         index=0 if example_data["major"] == "Công nghệ thông tin" else 1)
else:
    name = st.text_input("👤 Họ và tên")
    student_id = st.text_input("🎓 Mã sinh viên")
    major = st.selectbox("📚 Ngành học", ["Công nghệ thông tin", "Kinh tế"])

st.write("---")

# ===============================
#     INPUT FUNCTIONS
# ===============================
def input_semester(semester_label, default_values=None):
    with st.expander(f"📖 {semester_label}", expanded=True):
        col1, col2 = st.columns(2)

        with col1:
            somon0thi = st.number_input("Số môn không thi", 0,
                                        value=default_values[0] if default_values else 0,
                                        key=f"sm0_{semester_label}")

            sotc0thi = st.number_input("Số tín chỉ không thi", 0,
                                       value=default_values[1] if default_values else 0,
                                       key=f"tc0_{semester_label}")

            sotcno = st.number_input("Số tín chỉ nợ", 0,
                                     value=default_values[2] if default_values else 0,
                                     key=f"tcno_{semester_label}")

            mhno = st.number_input("Số môn không đạt", 0,
                                   value=default_values[3] if default_values else 0,
                                   key=f"mhno_{semester_label}")

        with col2:
            TCHK = st.number_input("Tổng tín chỉ học kỳ", 0,
                                   value=default_values[4] if default_values else 0,
                                   key=f"tchk_{semester_label}")

            smhk = st.number_input("Số môn học kỳ", 0,
                                   value=default_values[5] if default_values else 0,
                                   key=f"smhk_{semester_label}")

            mhpass = st.number_input("Số môn đạt", 0,
                                     value=default_values[6] if default_values else 0,
                                     key=f"mhpass_{semester_label}")

            TBCHK = st.number_input("GPA", 0.0, 10.0,
                                    value=default_values[7] if default_values else 0.0,
                                    step=0.01,
                                    key=f"gpa_{semester_label}")

            xep_loai_selected = st.selectbox("Xếp loại", list(range(7)),
                                             index=default_values[8] if default_values else 0,
                                             key=f"xeploai_{semester_label}")

        sotc_qua = st.number_input("Số tín chỉ qua môn", 0,
                                   value=TCHK - sotcno if default_values else 0,
                                   key=f"tcqua_{semester_label}")

    return [
        somon0thi, sotc0thi, sotcno, mhno,
        TCHK, smhk, mhpass, TBCHK, xep_loai_selected, sotc_qua
    ]


def input_important_features(semester_label, default_values=None):
    with st.expander(f"📘 {semester_label}", expanded=True):
        col1, col2 = st.columns(2)

        with col1:
            sotc_qua = st.number_input("Số tín chỉ đạt", 0,
                                       value=default_values[0] if default_values else 0,
                                       key=f"tcqua_imp_{semester_label}")

            sotcno = st.number_input("Số tín chỉ nợ", 0,
                                     value=default_values[1] if default_values else 0,
                                     key=f"tcno_imp_{semester_label}")

        with col2:
            TBCHK = st.number_input("Điểm trung bình", 0.0, 10.0,
                                    value=default_values[2] if default_values else 0.0,
                                    step=0.01,
                                    key=f"gpa_imp_{semester_label}")

    return [sotc_qua, sotcno, TBCHK]

# ===============================
#     GENERATE FINAL FEATURES
# ===============================
data = []
semesters = ["HỌC KỲ I", "HỌC KỲ II", "HỌC KỲ III", "HỌC KỲ IV", "HỌC KỲ V", "HỌC KỲ VI"]

model_type = st.sidebar.selectbox("🧠 Chọn mô hình dự báo:",
                                  ["Dùng toàn bộ đặc trưng", "Dùng đặc trưng quan trọng"])

# ===============================
#     FULL MODEL (61 features)
# ===============================
if model_type == "Dùng toàn bộ đặc trưng":
    st.subheader("🔢 Nhập thông tin học kỳ (Full Data)")

    full_features = []
    sotc_qua_list = []

    for idx, sem in enumerate(semesters):
        default_values = None
        if example_data and len(example_data["semester_data"]) == 54:
            default_values = example_data["semester_data"][idx*9:(idx+1)*9]

        sem_values = input_semester(sem, default_values)

        # 9 features đầu
        full_features += sem_values[:9]

        # sotc_qua để riêng
        sotc_qua_list.append(sem_values[9])

    nganh = 0 if major == "Công nghệ thông tin" else 1

    # THỨ TỰ MÔ HÌNH FULL:
    final_features = np.array(
        full_features + [nganh] + sotc_qua_list
    ).reshape(1, -1)

# ===============================
#     SIMPLE MODEL (18 features)
# ===============================
else:
    st.subheader("✨ Nhập thông tin rút gọn (Simple Data)")

    simple_features = []

    for idx, sem in enumerate(semesters):
        default_values = None
        if example_data and len(example_data["semester_data"]) == 18:
            default_values = example_data["semester_data"][idx*3:(idx+1)*3]

        sotc_qua, sotcno, TBCHK = input_important_features(sem, default_values)

        # đúng thứ tự mô hình simple
        simple_features += [sotc_qua, sotcno, TBCHK]

    final_features = np.array(simple_features).reshape(1, -1)
# ===============================
#         PREDICT BUTTON
# ===============================
st.write("---")

if st.button("🎯 DỰ BÁO"):
    # FULL FEATURES (61 inputs)
    if model_type == "Dùng toàn bộ đặc trưng":
        prediction = model_full.predict(final_features)

        if prediction[0] == 1:
            st.success(f"🎉 KẾT QUẢ: BẠN {name} CÓ KHẢ NĂNG TỐT NGHIỆP ĐÚNG HẠN!")
            st.balloons()
        else:
            st.error(f"⚠️ KẾT QUẢ: BẠN {name} CÓ NGUY CƠ **KHÔNG TỐT NGHIỆP ĐÚNG HẠN**.")
            st.snow()

    # SIMPLE FEATURES (18 inputs)
    else:
        prediction = model_important.predict(final_features)

        if prediction[0] == 1:
            st.success(f"🎉 KẾT QUẢ: BẠN {name} CÓ KHẢ NĂNG TỐT NGHIỆP ĐÚNG HẠN!")
            st.balloons()
        else:
            st.error(f"⚠️ KẾT QUẢ: BẠN {name} CÓ NGUY CƠ **KHÔNG TỐT NGHIỆP ĐÚNG HẠN**.")
            st.snow()