app.py

import streamlit as st
import pandas as pd
import xgboost as xgb
import os
import numpy as np

# --- CONFIGURATION ---
PAGE_TITLE = "FTC DECODE Match Predictor"
STATS_FILE = 'team_stats_split.csv'
MODEL_AUTO = 'model_auto.json'
MODEL_TELE = 'model_teleop.json'

st.set_page_config(page_title=PAGE_TITLE, layout="centered")


# --- 1. BACKEND LOGIC ---

@st.cache_resource
def load_resources():
    """Loads the models and stats once and keeps them in memory."""
    if not os.path.exists(MODEL_AUTO) or not os.path.exists(MODEL_TELE):
        return None, None, None
    if not os.path.exists(STATS_FILE):
        return None, None, None

    m_auto = xgb.XGBRegressor()
    m_auto.load_model(MODEL_AUTO)

    m_tele = xgb.XGBRegressor()
    m_tele.load_model(MODEL_TELE)

    # Load Stats
    stats = pd.read_csv(STATS_FILE).set_index('team')
    return m_auto, m_tele, stats


def get_prediction(m_auto, m_tele, stats, r1, r2, b1, b2, season_progress=1.0):
    feature_names = m_auto.get_booster().feature_names
    row = {'season_progress': season_progress}

    def fill(team, prefix):
        if team in stats.index:
            for k, v in stats.loc[team].items():
                if f"{prefix}_{k}" in feature_names:
                    row[f"{prefix}_{k}"] = v

    fill(r1, 'red_1');
    fill(r2, 'red_2')
    fill(b1, 'blue_1');
    fill(b2, 'blue_2')

    # Predict RED
    df_red = pd.DataFrame([row], columns=feature_names).fillna(0)
    ra = m_auto.predict(df_red)[0]
    rt = m_tele.predict(df_red)[0]

    # Predict BLUE (Flip inputs)
    row_blue = {'season_progress': season_progress}
    for col in feature_names:
        if 'red_1' in col:
            row_blue[col] = row.get(col.replace('red_1', 'blue_1'), 0)
        elif 'red_2' in col:
            row_blue[col] = row.get(col.replace('red_2', 'blue_2'), 0)
        elif 'blue_1' in col:
            row_blue[col] = row.get(col.replace('blue_1', 'red_1'), 0)
        elif 'blue_2' in col:
            row_blue[col] = row.get(col.replace('blue_2', 'red_2'), 0)

    df_blue = pd.DataFrame([row_blue], columns=feature_names).fillna(0)
    ba = m_auto.predict(df_blue)[0]
    bt = m_tele.predict(df_blue)[0]

    return (ra, rt), (ba, bt)


# --- 2. FRONTEND UI ---

st.title(f"🤖 {PAGE_TITLE}")
st.markdown("### AI-Powered Alliance Prediction")

m_auto, m_tele, stats = load_resources()

if m_auto is None:
    st.error("⚠️ Models or Stats file missing! Run `train_split.py` first.")
    st.stop()

# Sidebar
st.sidebar.header("Match Settings")
penalty_bonus = st.sidebar.slider("Est. Penalty Points", 0, 50, 15)
season_prog = st.sidebar.slider("Season Progress", 0.1, 1.0, 1.0)

# Inputs
col1, col2 = st.columns(2)
with col1:
    st.subheader("🔴 Red Alliance")
    r1 = st.number_input("Red Team 1", value=0, step=1, format="%d")
    r2 = st.number_input("Red Team 2", value=0, step=1, format="%d")
with col2:
    st.subheader("🔵 Blue Alliance")
    b1 = st.number_input("Blue Team 1", value=0, step=1, format="%d")
    b2 = st.number_input("Blue Team 2", value=0, step=1, format="%d")

# --- PREDICTION LOGIC ---
if st.button("🚀 Predict Match Result", type="primary"):
    # 1. Input Validation
    if r1 == 0 or r2 == 0 or b1 == 0 or b2 == 0:
        st.warning("Please enter valid team numbers for all slots.")
        st.stop()

    # 2. Database Check (The New Feature)
    teams = [r1, r2, b1, b2]
    missing_teams = [t for t in teams if t not in stats.index]

    if missing_teams:
        st.error(f"❌ Error: The following teams are not in the database: {', '.join(map(str, missing_teams))}")
        st.info("The model cannot predict matches for teams with no history. Please check the team numbers.")
    else:
        # 3. Run Prediction
        (ra, rt), (ba, bt) = get_prediction(m_auto, m_tele, stats, r1, r2, b1, b2, season_prog)

        red_total = ra + rt + penalty_bonus
        blue_total = ba + bt + penalty_bonus
        diff = abs(red_total - blue_total)

        winner = "RED WINS" if red_total > blue_total else "BLUE WINS"
        win_color = "red" if red_total > blue_total else "blue"

        st.divider()
        st.markdown(f"<h1 style='text-align: center; color: {win_color};'>{winner} (+{diff:.1f})</h1>",
                    unsafe_allow_html=True)

        res_col1, res_col2 = st.columns(2)
        with res_col1:
            st.markdown(f"### 🔴 Red: {red_total:.0f}")
            st.write(f"**Auto:** {ra:.1f} | **Teleop:** {rt:.1f}")
        with res_col2:
            st.markdown(f"### 🔵 Blue: {blue_total:.0f}")
            st.write(f"**Auto:** {ba:.1f} | **Teleop:** {bt:.1f}")

        st.divider()
        st.subheader("📊 Team Stats (OPR)")

        # Display Stats Table
        team_data = []
        for t in teams:
            row = stats.loc[t]
            team_data.append({
                "Team": t,
                "Auto OPR": f"{row.get('opr_auto_score', 0):.1f}",
                "Teleop OPR": f"{row.get('opr_dc_score', 0):.1f}",
                "Auto Artifacts": f"{row.get('opr_auto_artifacts', 0):.1f}",
                "Teleop Artifacts": f"{row.get('opr_teleop_artifacts', 0):.1f}",
            })
        st.dataframe(pd.DataFrame(team_data), hide_index=True)